Add supports_gradient_checkpointing to NemotronHPreTrainedModel (#45625)

Add supports_gradient_checkpointing to NemotronHPreTrainedModel

Author: sergiopaniego

src/transformers/models/nemotron_h/modeling_nemotron_h.py

@@ -952,6 +952,7 @@ def forward(
 class NemotronHPreTrainedModel(PreTrainedModel):
     config: NemotronHConfig
     base_model_prefix = "model"
+    supports_gradient_checkpointing = True
     _no_split_modules = ["NemotronHBlock"]
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True

src/transformers/models/nemotron_h/modular_nemotron_h.py

@@ -305,6 +305,7 @@ def forward(
 class NemotronHPreTrainedModel(PreTrainedModel):
     config: NemotronHConfig
     base_model_prefix = "model"
+    supports_gradient_checkpointing = True
     _no_split_modules = ["NemotronHBlock"]
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True