Update mllama_audio_model.py

mllama_audio_model.py

@@ -35,12 +35,11 @@ class MllamaAudioModel(MllamaPreTrainedModel):
         for i in range(bs):
             for j in range(max_num_img):
                 audio_id = -1 - j
-                idx = torch.where(input_ids[i] == audio_id)
-                if idx.numel() > 0:
-                    input_embeddings[i][idx] = torch.concat([self.start_of_audio, audio_features[i, j][idx], self.end_of_audio])
+                if torch.any(input_ids[i] == audio_id):
+                    idx = input_ids[i] == audio_id
+                    input_embeddings[i][idx] = torch.concat([self.start_of_audio, audio_features[i, j], self.end_of_audio])
 
-        idx = torch.where(input_ids < 0 and input_ids >= -max_num_img)
-        input_ids[idx].fill_(self.filler_token_id)
+        input_ids[input_ids < 0].fill_(self.filler_token_id)
 
         if return_dict:
             return dict(input_embeddings=input_embeddings)