Update utils.py

utils.py

@@ -191,7 +191,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")
         )
@@ -284,7 +284,10 @@ class MolecularPropertyPredictionModel():
             for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
                 with torch.no_grad():
                     batch = {k: v.to(self.base_model.device) for k, v in batch.items()}
+                    print(self.base_model.device)
+                    print(batch)
                     outputs = self.base_model(**batch)
+                    print(output)
                 if task_type == "regression": # TODO: check if the model is regression or classification
                     y_pred.append(outputs.logits.cpu().detach().numpy())
                 else: