Update modeling_aimv2.py

modeling_aimv2.py

@@ -222,7 +222,7 @@ class AIMv2Model(AIMv2PretrainedModel):
             hidden_states=hidden_states,
         )
 
-
+'''
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
         super().__init__(config)
@@ -310,9 +310,9 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             hidden_states=outputs.hidden_states,
             # attentions=outputs.attentions,
         )
+'''
 
 
-'''
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
         super().__init__(config)
@@ -338,15 +338,10 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[tuple, ImageClassifierOutput]:
-        print("Forward pass initiated")
-        print(f"Input pixel_values shape: {pixel_values.shape if pixel_values is not None else 'None'}")
-        print(f"Head mask provided: {head_mask is not None}")
-        print(f"Labels provided: {labels is not None}")
         
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
-        print(f"Using return_dict: {return_dict}")
 
         # Call base model
         outputs = self.aimv2(
@@ -356,33 +351,23 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             return_dict=return_dict,
         )
         sequence_output = outputs[0]
-        print(f"Sequence output shape: {sequence_output.shape}")
-
         # Classifier head
         logits = self.classifier(sequence_output[:, 0, :])
-        print(f"Logits shape: {logits.shape}")
-        print(f"Logits shape: {logits}")
 
         loss = None
         if labels is not None:
             labels = labels.to(logits.device)
-            print(f"Labels shape: {labels.shape}")
-            print(f"Labels shape: {labels}")
-            
             # Always use cross-entropy loss
             loss_fct = CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
-            print(f"Loss computed: {loss.item()}")
 
         if not return_dict:
             output = (logits,) + outputs[1:]
-            print("Returning as tuple")
             return ((loss,) + output) if loss is not None else output
 
-        print("Returning as ImageClassifierOutput")
         return ImageClassifierOutput(
             loss=loss,
             logits=logits,
             hidden_states=outputs.hidden_states,
         )
-'''
+