Update prediction.py

prediction.py

@@ -17,7 +17,8 @@ np.random.seed(42)
 torch.manual_seed(42)
 torch.backends.cudnn.deterministic = True
 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 ------------------------
 @st.cache_resource
 def load_chemberta():
@@ -26,8 +27,7 @@ def load_chemberta():
     model.eval()
     model.to(device)  # Send model to GPU if available
     return tokenizer, model
-    # Check if CUDA is available for GPU acceleration
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 # ------------------------ Load Scalers ------------------------
 scalers = {
@@ -66,7 +66,6 @@ class TransformerRegressor(nn.Module):
         x = x.mean(dim=1)
         return self.regression_head(x)
 
-# ------------------------ Load Model ------------------------
 @st.cache_resource
 def load_model():
     # Initialize the model architecture first