Add supports_gradient_checkpointing

2025-03-26 13:52:52 +08:00
parent d13495871e
commit b2bedcd779
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']