[FIX] weight_update

model.py

@@ -44,7 +44,7 @@ class Architecture(nn.Module):
 
 def load_model():
     model = Architecture()
-    model.load_state_dict(torch.load('model_weights.pth'))
+    model.load_state_dict(torch.load('./model_weights.pth'))
     return model
 
 def inference_model(model, input):

model_hugging_face.py

@@ -1,3 +1,4 @@
+from transformers import AutoConfig, AutoModel
 from transformers import PreTrainedModel, PretrainedConfig
 import torch.nn as nn
 import torch
@@ -59,7 +60,11 @@ class Architecture(PreTrainedModel):
 
 # Loading the model from saved weights
 def load_model():
+    AutoConfig.register("architecture", ArchitectureConfig)
+    AutoModel.register(ArchitectureConfig, Architecture)
     config = ArchitectureConfig()
     model = Architecture(config)
-    model.load_state_dict(torch.load('model_weights.pth'))
+    model.load_state_dict(torch.load('./model_weights.pth'))
     return model
+
+load_model()