Update modeling_aimv2.py

modeling_aimv2.py

@@ -311,8 +311,9 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
 
 import logging
 
+
 # Setup logging
-logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
@@ -340,15 +341,15 @@ 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}")
+        logger.info("Forward pass initiated")
+        logger.info(f"Input pixel_values shape: {pixel_values.shape if pixel_values is not None else 'None'}")
+        logger.info(f"Head mask provided: {head_mask is not None}")
+        logger.info(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}")
+        logger.info(f"Using return_dict: {return_dict}")
 
         # Call base model
         outputs = self.aimv2(
@@ -358,28 +359,28 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             return_dict=return_dict,
         )
         sequence_output = outputs[0]
-        logger.debug(f"Sequence output shape: {sequence_output.shape}")
+        logger.info(f"Sequence output shape: {sequence_output.shape}")
 
         # Classifier head
         logits = self.classifier(sequence_output[:, 0, :])
-        logger.debug(f"Logits shape: {logits.shape}")
+        logger.info(f"Logits shape: {logits.shape}")
 
         loss = None
         if labels is not None:
             labels = labels.to(logits.device)
-            logger.debug(f"Labels shape: {labels.shape}")
+            logger.info(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()}")
+            logger.info(f"Loss computed: {loss.item()}")
 
         if not return_dict:
             output = (logits,) + outputs[1:]
-            logger.debug("Returning as tuple")
+            logger.info("Returning as tuple")
             return ((loss,) + output) if loss is not None else output
 
-        logger.debug("Returning as ImageClassifierOutput")
+        logger.info("Returning as ImageClassifierOutput")
         return ImageClassifierOutput(
             loss=loss,
             logits=logits,