propagate device setting to calculator

atomind-ai · Jan 14, 2025 · 8cb1d3b · 8cb1d3b
1 parent a909029
commit 8cb1d3b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/mlip_arena/tasks/utils.py b/mlip_arena/tasks/utils.py
@@ -28,11 +28,13 @@ def get_calculator(
     device: str | None = None,
 ) -> Calculator | SumCalculator:
     """Get a calculator with optional dispersion correction."""
-    device = device or str(get_freer_device())
 
-    logger.info(f"Using device: {device}")
+    device = device or str(get_freer_device())
 
     calculator_kwargs = calculator_kwargs or {}
+    calculator_kwargs.update({"device": device})
+
+    logger.info(f"Using device: {device}")
 
     if isinstance(calculator_name, MLIPEnum) and calculator_name in MLIPEnum:
         calc = calculator_name.value(**calculator_kwargs)