Update prediction.py

prediction.py

@@ -68,12 +68,9 @@ class TransformerRegressor(nn.Module):
 # ------------------------ Load Model ------------------------
 @st.cache_resource
 def load_model():
-    model = TransformerRegressor()
-    state_dict = torch.load("transformer_model.bin", map_location=torch.device("cpu"))
-    model.load_state_dict(state_dict)
+    model = torch.load("transformer_model.bin", map_location=torch.device("cpu"))
     model.eval()
     return model
-
 model = load_model()
 
 # ------------------------ Descriptors ------------------------