device

main.py

@@ -42,10 +42,10 @@ atom_mapping = {0:'H', 1:'C', 2:'N', 3:'O', 4:'F', 5:'P', 6:'S', 7:'CL', 8:'BR',
 model = GNN_MD(11, 64)
 state_dict = torch.load(
     "best_weights_rep0.pt",
-    map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    map_location=torch.device("cpu"),
 )["model_state_dict"]
 model.load_state_dict(state_dict)
-model = model.to("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to('cpu')
 model.eval()