Update audio_text_multimodal.py

audio_text_multimodal.py

@@ -12,6 +12,14 @@ from transformers import (
     Wav2Vec2Model
 )
 
+from transformers.models.wavlm.modeling_wavlm import (
+    WavLMEncoder,
+    WavLMEncoderStableLayerNorm,
+    WavLMFeatureEncoder
+)
+
+from transformers.models.bert.modeling_bert import BertEncoder
+
 
 class MultiModalConfig(PretrainedConfig):
     """Base class for multimodal configs"""
@@ -191,6 +199,8 @@ class WavLMBertForSequenceClassification(AudioTextFusionModelForSequenceClassifi
     """
     def __init__(self, config):
         super().__init__(config)
+        self.supports_gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+        
         self.audio_config = WavLMConfig.from_dict(self.config.WavLMModel)
         self.text_config = BertConfig.from_dict(self.config.BertModel)
         self.audio_model = WavLMModel(self.audio_config)
@@ -215,6 +225,11 @@ class WavLMBertForSequenceClassification(AudioTextFusionModelForSequenceClassifi
             (cls_dim * 2) // self.config.kernel_size, self.config.num_labels
         )
         self.init_weights()
+
+    @staticmethod
+    def _set_gradient_checkpointing(module, value=False):
+        if isinstance(module, (WavLMEncoder, WavLMEncoderStableLayerNorm, WavLMFeatureEncoder, BertEncoder)):
+            module.gradient_checkpointing = value
         
     def forward(
             self,