Update prediction.py

prediction.py

@@ -70,19 +70,16 @@ class TransformerRegressor(nn.Module):
 
 @st.cache_resource
 def load_model():
-    # Initialize the model architecture first
     model = TransformerRegressor()
-
-    # Load the state_dict (weights) from the saved model file
-    state_dict = torch.load("transformer_model(1).bin", map_location=device)  # Ensure loading on the correct device
-
-    # Load the state_dict into the model
+    state_dict = torch.load("transformer_model(1).bin", map_location=device)
     model.load_state_dict(state_dict)
-
-    # Set the model to evaluation mode
     model.eval()
-    model.to(device)  # Send model to GPU if available
+    model.to(device)
     return model
+
+# Call them to load the actual models
+tokenizer, chemberta = load_chemberta()
+model = load_model()
 # ------------------------ Descriptors ------------------------
 def compute_descriptors(smiles: str):
     mol = Chem.MolFromSmiles(smiles)