registering

app.py

@@ -11,7 +11,7 @@ import sys
 # import timm
 from types import SimpleNamespace
 # from transformers import AutoModel, pipeline
-from transformers import AutoModelForImageClassification, AutoModel
+from transformers import AutoModelForImageClassification, AutoModel, AutoConfig
 import torch
 
 sys.path.insert(1, "../")
@@ -254,6 +254,17 @@ def predict_and_analyze(model_name, num_channels, dim, input_channel, image):
         EfficientNetConfig.model_type = hparams.model_type
     
     config.save_pretrained(save_directory=model_loading_name)
+    
+    model = EfficientNet(dropout=hparams.dropout,
+                                    num_channels=hparams.num_channels,
+                                    num_classes=hparams.num_classes,
+                                    size=hparams.size,
+                                    stochastic_depth_prob=hparams.stochastic_depth_prob,
+                                    width_mult=hparams.width_mult,
+                                    depth_mult=hparams.depth_mult,)
+    
+    AutoConfig.register(model_loading_name, config)
+    AutoModel.register(config, model)
     # config = EfficientNetConfig.from_pretrained(model_loading_name)
     
     # model = EfficientNetPreTrained.from_pretrained(model_loading_name)