Add supports_gradient_checkpointing

2025-03-26 15:35:47 +08:00
parent ee45ba3769
commit d0b22d04cd
3 changed files with 14 additions and 0 deletions
--- a/modeling_intern_vit.py
+++ b/modeling_intern_vit.py
@@ -364,6 +364,7 @@ class InternVisionEncoder(nn.Module):
 class InternVisionModel(PreTrainedModel):
    main_input_name = 'pixel_values'
    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
    config_class = InternVisionConfig
    _no_split_modules = ['InternVisionEncoderLayer']