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)
@@ -306,3 +306,64 @@ class AIMv2ForImageClassification(AIMv2PretrainedModel):
             hidden_states=outputs.hidden_states,
             # attentions=outputs.attentions,
         )
+'''
+
+
+class AIMv2ForImageClassification(AIMv2PretrainedModel):
+    def __init__(self, config: AIMv2Config):
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.aimv2 = AIMv2Model(config)
+
+        # Classifier head
+        self.classifier = (
+            nn.Linear(config.hidden_size, config.num_labels)
+            if config.num_labels > 0
+            else nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[tuple, ImageClassifierOutput]:
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.aimv2(
+            pixel_values,
+            mask=head_mask,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.classifier(sequence_output[:, 0, :])
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            
+            # Always use cross-entropy loss
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return ImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+        )