Update prediction.py

prediction.py

@@ -20,13 +20,15 @@ torch.backends.cudnn.benchmark = False
 # Check if CUDA is available for GPU acceleration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # ------------------------ Load ChemBERTa Model + Tokenizer ------------------------
+# ------------------------ Load ChemBERTa Model + Tokenizer ------------------------
 @st.cache_resource
 def load_chemberta():
     tokenizer = AutoTokenizer.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
     model = AutoModel.from_pretrained("seyonec/ChemBERTa-zinc-base-v1")
     model.eval()
-    model.to(device)  # Send model to GPU if available
+    model.to(device)
     return tokenizer, model
+tokenizer, chemberta = load_chemberta()
 
 
 # ------------------------ Load Scalers ------------------------