this sucks

app.py

@@ -10,6 +10,7 @@ import sys
 import timm
 from types import SimpleNamespace
 # from transformers import AutoModel, pipeline
+from transformers import AutoModelForImageClassification
 import torch
 
 sys.path.insert(1, "../")
@@ -176,7 +177,6 @@ def plot_activations(activation_1: list, activation_2: list, origin='lower'):
     
     return fig
 
-
 def predict_and_analyze(model_name, num_channels, dim, input_channel, image):
     
     '''
@@ -203,7 +203,7 @@ def predict_and_analyze(model_name, num_channels, dim, input_channel, image):
     print("Data loaded")
     
     print("Loading model")
-    model_loading_name =  model_path + "%s_%i_planet_detection" % (model_name, num_channels)
+    model_loading_name = "%s_%i_planet_detection" % (model_name, num_channels)
     
     if 'eff' in model_name:
         hparams = effnet_hparams[num_channels]
@@ -220,13 +220,15 @@ def predict_and_analyze(model_name, num_channels, dim, input_channel, image):
                                     depth_mult=hparams.depth_mult,
         )
     
-    config.save_pretrained("%s_planet_detection" % (model_name))
-    
-    config = EfficientNetConfig.from_pretrained("%s_planet_detection" % (model_name))
+    config.save_pretrained(model_loading_name)
+    config = EfficientNetConfig.from_pretrained(model_loading_name)
     
     model = EfficientNetPreTrained(config)
+    
+    # config.register_for_auto_class()
+    # model.register_for_auto_class("AutoModelForImageClassification")
 
-    pretrained_model = timm.create_model("%s_planet_detection" % (model_name), pretrained=True)
+    pretrained_model = timm.create_model(model_loading_name, pretrained=True)
     model.model.load_state_dict(pretrained_model.state_dict())
     
     # pipeline = pipeline(task="image-classification", model=model_loading_name)