Upload model

config.json

@@ -1,12 +1,10 @@
 {
-  "_name_or_path": "./efficientnet/temp",
   "architectures": [
-    "EfficientNetModelForImageClassification"
+    "EfficientNetModel"
   ],
   "auto_map": {
     "AutoConfig": "configuration_efficientnet.EfficientNetConfig",
-    "AutoModel": "modeling_efficientnet.EfficientNetModel",
-    "AutoModelForImageClassification": "modeling_efficientnet.EfficientNetModelForImageClassification"
+    "AutoModel": "modeling_efficientnet.EfficientNetModel"
   },
   "model_name": "efficientnet_b0",
   "model_type": "efficientnet",

modeling_efficientnet.py

@@ -2,6 +2,7 @@ from torch import nn
 from functools import partial
 
 from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
 from timm import create_model
 
 from configuration_efficientnet import EfficientNetConfig
@@ -15,7 +16,10 @@ class EfficientNetModel(PreTrainedModel):
         self.model = create_model(config.model_name, pretrained = config.pretrained)
     
     def forward(self, pixel_values):
-        return self.model.forward_features(pixel_values)
+        last_hidden_state = self.model.forward_features(pixel_values)
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state = last_hidden_state
+        )
 
 class EfficientNetModelForImageClassification(PreTrainedModel):
     config_class = EfficientNetConfig
@@ -29,8 +33,10 @@ class EfficientNetModelForImageClassification(PreTrainedModel):
         logits = self.model(pixel_values)
         if labels is not None:
             loss = nn.CrossEntropyLoss(logits, labels)
-            return {"loss": loss, "logits": logits}
-        return logits
+        return ImageClassifierOutputWithNoAttention(
+            loss = loss,
+            logits = logits
+        )
     
 __all__ = [
     "EfficientNetModel",