Update modeling_aimv2.py

modeling_aimv2.py

@@ -222,7 +222,7 @@ class AIMv2Model(AIMv2PretrainedModel):
             hidden_states=hidden_states,
         )
 
-'''
+
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
         super().__init__(config)
@@ -271,6 +271,7 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
 
         loss = None
         if labels is not None:
+            print("LABELS: ", labels)
             # move labels to correct device to enable model parallelism
             labels = labels.to(logits.device)
             if self.config.problem_type is None:
@@ -295,7 +296,10 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             elif self.config.problem_type == "multi_label_classification":
                 loss_fct = BCEWithLogitsLoss()
                 loss = loss_fct(logits, labels)
-
+        
+        print("PROBLEM", self.config.problem_type)
+        print("LOSS: ", loss)
+        
         if not return_dict:
             output = (logits,) + outputs[1:]
             return ((loss,) + output) if loss is not None else output
@@ -306,9 +310,9 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             hidden_states=outputs.hidden_states,
             # attentions=outputs.attentions,
         )
-'''
 
 
+'''
 class AIMv2ForImageClassification(AIMv2PretrainedModel):
     def __init__(self, config: AIMv2Config):
         super().__init__(config)
@@ -381,4 +385,4 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             logits=logits,
             hidden_states=outputs.hidden_states,
         )
-
+'''