Update modeling_aimv2.py

modeling_aimv2.py

@@ -357,11 +357,13 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
         # 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()