Update prediction.py

prediction.py

@@ -69,7 +69,8 @@ class TransformerRegressor(nn.Module):
 @st.cache_resource
 def load_model():
     model = TransformerRegressor()
-    model.load_state_dict(torch.load("transformer_model.pt", map_location=torch.device("cpu")))
+    state_dict = torch.load("transformer_model.bin", map_location=torch.device("cpu"))
+    model.load_state_dict(state_dict)
     model.eval()
     return model