Update prediction.py

prediction.py

@@ -65,12 +65,24 @@ class TransformerRegressor(nn.Module):
         x = x.mean(dim=1)
         return self.regression_head(x)
 
+# ------------------------ Load Model ------------------------
 # ------------------------ Load Model ------------------------
 @st.cache_resource
 def load_model():
-    model = torch.load("transformer_model.bin", map_location=torch.device("cpu"))
+    # Initialize the model architecture first
+    model = TransformerRegressor()
+
+    # Load the state_dict (weights) from the saved model file
+    state_dict = torch.load("transformer_model.bin", map_location=torch.device("cpu"))
+
+    # Load the state_dict into the model
+    model.load_state_dict(state_dict)
+
+    # Set the model to evaluation mode
     model.eval()
     return model
+
+# Load the model
 model = load_model()
 
 # ------------------------ Descriptors ------------------------