Fixed device choice in mtt eval

metatensor · Nov 6, 2024 · add538d · add538d
1 parent 5a60116
commit add538d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/metatrain/cli/eval.py b/src/metatrain/cli/eval.py
@@ -184,7 +184,9 @@ def _eval_targets(
     # Infer the device and dtype from the model
     model_tensor = next(itertools.chain(model.parameters(), model.buffers()))
     dtype = model_tensor.dtype
-    device = "cuda" if "cuda" in model.capabilities().supported_devices else "cpu"
+    device = "cpu"
+    if torch.cuda.is_available() and "cuda" in model.capabilities().supported_devices:
+        device = "cuda"
     logger.info(f"Running on device {device} with dtype {dtype}")
     model.to(dtype=dtype, device=device)