Update modeling_aimv2.py

modeling_aimv2.py

@@ -315,10 +315,14 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
 
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
+        print("Initializing AIMv2ForImageClassification")
         super().__init__(config)
 
         self.num_labels = config.num_labels
+        print(f"Number of labels: {self.num_labels}")
+        
         self.aimv2 = AIMv2Model(config)
+        print("Initialized AIMv2 base model")
 
         # Classifier head
         self.classifier = (
@@ -326,9 +330,11 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             if config.num_labels > 0
             else nn.Identity()
         )
+        print(f"Initialized classifier: {self.classifier}")
 
         # Initialize weights and apply final processing
         self.post_init()
+        print("Weights initialized and final processing applied")
 
     def forward(
         self,
@@ -338,33 +344,48 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[tuple, ImageClassifierOutput]:
+        print("Forward pass started")
         
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
+        print(f"return_dict: {return_dict}")
 
         # Call base model
+        print("Calling AIMv2 base model")
         outputs = self.aimv2(
             pixel_values,
             mask=head_mask,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
+        print(f"AIMv2 outputs received: {outputs}")
+        
         sequence_output = outputs[0]
+        print(f"Shape of sequence_output: {sequence_output.shape}")
+
         # Classifier head
         logits = self.classifier(sequence_output[:, 0, :])
+        print(f"Logits calculated: {logits.shape}")
 
         loss = None
         if labels is not None:
+            print(labels)
+            print(f"Labels provided: {labels.shape}")
             labels = labels.to(logits.device)
+            print(f"Labels moved to device: {labels.device}")
+            
             # Always use cross-entropy loss
             loss_fct = CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            print(f"Loss calculated: {loss.item()}")
 
         if not return_dict:
             output = (logits,) + outputs[1:]
+            print("Output without return_dict")
             return ((loss,) + output) if loss is not None else output
 
+        print("Returning ImageClassifierOutput")
         return ImageClassifierOutput(
             loss=loss,
             logits=logits,