Update any_model.py

any_model.py

@@ -667,11 +667,7 @@ class AnyModelForConditionalGeneration(AnyModelPreTrainedModel):
             # 2. Merge text and images
             if pixel_values_1 is not None and pixel_values_1 is not None and input_ids.shape[1] != 1:
                 assert modality is not None, "modality must be provided when pixel_values is not None"
-                '''
-                if isinstance(modality, list):
-                    assert len(set(modality)) == 1, "only one kind modality can be provided in a batch"
-                    modality = modality[0]
-                '''
+ 
                 for i in range(2):
                     pixel_values = pixel_values_1 if i == 0 else pixel_values_2
                     if modality[0][i] == ModalityType.IMAGE:
@@ -698,12 +694,6 @@ class AnyModelForConditionalGeneration(AnyModelPreTrainedModel):
                         raise ValueError(f"modality {modality[i]} is not supported")
                     
                     inputs_embeds = inputs_embeds.to(features.dtype)
-                    '''
-                    print('+++'*10)
-                    print(input_ids)
-                    print(torch.sum(input_ids == self.config.audio_token_index, dim=-1))
-                    print('+++'*10)
-                    '''
                     inputs_embeds, attention_mask, labels, position_ids = self.merge_input_ids_with_other_features(
                         features, inputs_embeds, input_ids, attention_mask, labels
                     )
@@ -832,7 +822,7 @@ class AnyRewardModel(AnyModelForConditionalGeneration):
         attention_mask: torch.Tensor | None = None,
         **kwargs,
     ) -> torch.Tensor:
-        outputs = self.model(
+        outputs = super().forward(
             input_ids,
             attention_mask=attention_mask,
             output_hidden_states=True,