Update modeling_aimv2.py

modeling_aimv2.py

@@ -309,6 +309,12 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
 '''
 
 
+import logging
+
+# Setup logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
         super().__init__(config)
@@ -334,34 +340,46 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[tuple, ImageClassifierOutput]:
-
+        logger.debug("Forward pass initiated")
+        logger.debug(f"Input pixel_values shape: {pixel_values.shape if pixel_values is not None else 'None'}")
+        logger.debug(f"Head mask provided: {head_mask is not None}")
+        logger.debug(f"Labels provided: {labels is not None}")
+        
         return_dict = (
             return_dict if return_dict is not None else self.config.use_return_dict
         )
+        logger.debug(f"Using return_dict: {return_dict}")
 
+        # Call base model
         outputs = self.aimv2(
             pixel_values,
             mask=head_mask,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-
         sequence_output = outputs[0]
+        logger.debug(f"Sequence output shape: {sequence_output.shape}")
 
+        # Classifier head
         logits = self.classifier(sequence_output[:, 0, :])
+        logger.debug(f"Logits shape: {logits.shape}")
 
         loss = None
         if labels is not None:
             labels = labels.to(logits.device)
+            logger.debug(f"Labels shape: {labels.shape}")
             
             # Always use cross-entropy loss
             loss_fct = CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            logger.debug(f"Loss computed: {loss.item()}")
 
         if not return_dict:
             output = (logits,) + outputs[1:]
+            logger.debug("Returning as tuple")
             return ((loss,) + output) if loss is not None else output
 
+        logger.debug("Returning as ImageClassifierOutput")
         return ImageClassifierOutput(
             loss=loss,
             logits=logits,