going back to automodel

app.py

@@ -10,7 +10,7 @@ import sys
 # import timm
 from types import SimpleNamespace
 # from transformers import AutoModel, pipeline
-from transformers import AutoModelForImageClassification
+from transformers import AutoModelForImageClassification, AutoModel
 import torch
 
 sys.path.insert(1, "../")
@@ -220,12 +220,13 @@ def predict_and_analyze(model_name, num_channels, dim, input_channel, image):
                                     width_mult=hparams.width_mult,
                                     depth_mult=hparams.depth_mult,
         )
+        EfficientNetConfig.model_type = "efficientnet_%s_planet_detection" % (hparams.num_channels)
     
     config.save_pretrained(save_directory=model_loading_name)
     # config = EfficientNetConfig.from_pretrained(model_loading_name)
     
-    model = EfficientNetPreTrained.from_pretrained(model_loading_name)
-    # model = AutoModelForImageClassification.from_pretrained(model_loading_name, trust_remote_code=True)
+    # model = EfficientNetPreTrained.from_pretrained(model_loading_name)
+    model = AutoModel.from_pretrained(model_loading_name, trust_remote_code=True)
     
     # model = EfficientNetPreTrained(config)
     # config.register_for_auto_class()

model_utils/efficientnet_config.py

@@ -242,7 +242,7 @@ class FusedMBConv(nn.Module):
 
 class EfficientNetConfig(PretrainedConfig):
     
-    model_type = "efficientnet_61_planet_detection"
+    model_type = "efficientnet"
     
     def __init__(
         self,