Update hf_wrapper.py

hf_wrapper.py

@@ -1159,6 +1159,14 @@ class Effb2TrmCaptioningModel(PreTrainedModel):
         model = TransformerModel(encoder, decoder)
         self.model = ContraEncoderKdWrapper(model, config.shared_dim, config.tchr_dim)
     
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        model = super().from_pretrained(
+            pretrained_model_name_or_path, *args, **kwargs
+        )
+        model.model.model.decoder.word_embedding.weight = model.model.model.decoder.classifier.weight
+        return model
+        
     def forward(self,
                 audio: torch.Tensor,
                 audio_length: Union[List, np.ndarray, torch.Tensor],