Update audio_text_multimodal.py

audio_text_multimodal.py

@@ -14,6 +14,14 @@ from transformers import (
     Wav2Vec2Model
 )
 
+from transformers.models.wav2vec2.modeling_wav2vec2 import (
+    Wav2Vec2Encoder,
+    Wav2Vec2EncoderStableLayerNorm,
+    Wav2Vec2FeatureEncoder
+)
+
+from transformers.models.bert.modeling_bert import BertEncoder
+
 
 class MultiModalConfig(PretrainedConfig):
     """Base class for multimodal configs"""
@@ -170,7 +178,7 @@ class AudioTextModelForSequenceBaseClassification(BaseClassificationModel):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-        audio_mean = self.merged_strategy(audio_output.last_hidden_state, mode="mean")
+        audio_mean = self.merged_strategy(audio_output.last_hidden_state, mode=self.config.pooling_mode)
 
         pooled_output = torch.cat(
             (audio_mean, text_output.pooler_output), dim=1
@@ -205,6 +213,8 @@ class Wav2Vec2BertForSequenceClassification(AudioTextModelForSequenceBaseClassif
     """
     def __init__(self, config):
         super().__init__(config)
+        self.supports_gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+
         self.audio_config = Wav2Vec2Config.from_dict(self.config.Wav2Vec2Model)
         self.text_config = BertConfig.from_dict(self.config.BertModel)
         self.audio_model = Wav2Vec2Model(self.audio_config)
@@ -212,4 +222,9 @@ class Wav2Vec2BertForSequenceClassification(AudioTextModelForSequenceBaseClassif
         self.classifier = torch.nn.Linear(
             self.audio_config.hidden_size + self.text_config.hidden_size, self.num_labels
         )
-        self.init_weights()
+        self.init_weights()
+
+    @staticmethod
+    def _set_gradient_checkpointing(module, value=False):
+        if isinstance(module, (Wav2Vec2Encoder, Wav2Vec2EncoderStableLayerNorm, Wav2Vec2FeatureEncoder, BertEncoder)):
+            module.gradient_checkpointing = value