Update mllama_audio_model.py

mllama_audio_model.py

@@ -14,7 +14,7 @@ class Llama3Embedding(Wav2Vec2BertPreTrainedModel):
         assert config.output_hidden_size == text_config.hidden_size
         self.text_embeddings = nn.Embedding(text_config.vocab_size, text_config.hidden_size, text_config.pad_token_id)
         self.audio_embedding = Wav2Vec2BertModel(config)
-        assert self.text_embeddings.weight.shape[-1] == text_config.hidden_size, f"{self.text_embeddings.weight.shape=}, {text_config.hidden_size=}"
+        assert self.text_embeddings.data.weight.size(-1) == text_config.hidden_size, f"{self.text_embeddings.weight}, {text_config.hidden_size=}, {text_config.vocab_size=}"
         self.start_of_audio = nn.Parameter(data=torch.zeros((1, config.output_hidden_size)), requires_grad=True)
         self.end_of_audio = nn.Parameter(data=torch.zeros((1, config.output_hidden_size)), requires_grad=True)
         self.text_config = text_config