Update modeling_aimv2.py

modeling_aimv2.py

@@ -309,13 +309,6 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
 '''
 
 
-import logging
-
-
-# Setup logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
         super().__init__(config)
@@ -341,15 +334,15 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[tuple, ImageClassifierOutput]:
-        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}")
+        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
         )
-        logger.info(f"Using return_dict: {return_dict}")
+        print(f"Using return_dict: {return_dict}")
 
         # Call base model
         outputs = self.aimv2(
@@ -359,30 +352,31 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             return_dict=return_dict,
         )
         sequence_output = outputs[0]
-        logger.info(f"Sequence output shape: {sequence_output.shape}")
+        print(f"Sequence output shape: {sequence_output.shape}")
 
         # Classifier head
         logits = self.classifier(sequence_output[:, 0, :])
-        logger.info(f"Logits shape: {logits.shape}")
+        print(f"Logits shape: {logits.shape}")
 
         loss = None
         if labels is not None:
             labels = labels.to(logits.device)
-            logger.info(f"Labels shape: {labels.shape}")
+            print(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.info(f"Loss computed: {loss.item()}")
+            print(f"Loss computed: {loss.item()}")
 
         if not return_dict:
             output = (logits,) + outputs[1:]
-            logger.info("Returning as tuple")
+            print("Returning as tuple")
             return ((loss,) + output) if loss is not None else output
 
-        logger.info("Returning as ImageClassifierOutput")
+        print("Returning as ImageClassifierOutput")
         return ImageClassifierOutput(
             loss=loss,
             logits=logits,
             hidden_states=outputs.hidden_states,
         )
+