move to zero gpu

utils.py

@@ -14,7 +14,7 @@ import os
 import pickle
 from sklearn import preprocessing
 import json
-
+import spaces
 
 from rdkit import RDLogger, Chem
 # Suppress RDKit INFO messages
@@ -166,7 +166,7 @@ class MolecularPropertyPredictionModel():
         self.base_model = AutoModelForSequenceClassification.from_pretrained(
             "ChemFM/ChemFM-3B",
             config=config,
-            device_map="cpu",
+            device_map="cuda",
             trust_remote_code=True,
             token = os.environ.get("TOKEN")
         )
@@ -217,6 +217,7 @@ class MolecularPropertyPredictionModel():
             # handle error
             return "error"
     
+    @spaces.GPU
     def predict(self, valid_df, task_type):
         test_dataset = Dataset.from_pandas(valid_df)
         # construct the dataloader