make sure model evaluation runs on GPU

molmod · Jul 19, 2024 · c1f582f · c1f582f
1 parent 496caee
commit c1f582f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/psiflow/functions.py b/psiflow/functions.py
@@ -218,6 +218,8 @@ def __post_init__(self):
         from mace.tools import torch_tools, utils
 
         torch_tools.set_default_dtype(self.dtype)
+        if self.device == 'gpu':  # when it's not a specific GPU, use any
+            self.device = 'cuda'
         self.device = torch_tools.init_device(self.device)
 
         torch.set_num_threads(self.ncores)