Upload model

config.json

@@ -1,11 +1,10 @@
 {
   "architectures": [
-    "EfficientNetModelForImageClassification"
+    "EfficientNetModel"
   ],
   "auto_map": {
     "AutoConfig": "configuration_efficientnet.EfficientNetConfig",
-    "AutoModel": "modeling_efficientnet.EfficientNetModel",
-    "AutoModelForImageClassification": "modeling_efficientnet.EfficientNetModelForImageClassification"
+    "AutoModel": "modeling_efficientnet.EfficientNetModel"
   },
   "global_pool": "avg",
   "model_name": "efficientnet_b0",
@@ -13,5 +12,5 @@
   "num_classes": 1000,
   "pretrained": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.49.0"
+  "transformers_version": "4.50.3"
 }

modeling_efficientnet.py

@@ -85,11 +85,22 @@ class EfficientNetModelForImageClassification(PreTrainedModel):
             num_classes = config.num_classes, 
             global_pool = config.global_pool,
         )
+        
+    def compute_loss(self, logits, labels=None):
+        loss = None
+        if labels is None:
+            pass
+        else:
+            labels = tensor(labels)
+            ce_loss = nn.CrossEntropyLoss()
+            loss = ce_loss(logits, labels)
+
+        return loss
     
     def forward(
         self, 
         pixel_values: Tensor, 
-        labels: Optional[Union[List[int], Tensor]] = None
+        labels: Optional[Union[List[int], Tensor]] = None,
     ) -> ImageClassifierOutputWithNoAttention:
         """
         Parameters
@@ -105,21 +116,16 @@ class EfficientNetModelForImageClassification(PreTrainedModel):
             ImageClassifierOutputWithNoAttention
                 Object containing `logits` and `loss`.
         """
-        self.model.training = False if labels is None else True
-
         logits = self.model(pixel_values)
 
-        loss = None
-        if self.model.training:
-            labels = tensor(labels)
-            ce_loss = nn.CrossEntropyLoss()
-            loss = ce_loss(logits, labels)
+        loss = self.compute_loss(logits, labels)
 
         return ImageClassifierOutputWithNoAttention(
             loss = loss,
             logits = logits,
         )
     
+
 __all__ = [
     "EfficientNetModel",
     "EfficientNetModelForImageClassification"