fixed use_device in calculator init

CederGroupHub · Jul 24, 2024 · 6f7b035 · 6f7b035
1 parent 9556c3e
commit 6f7b035
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/chgnet/model/dynamics.py b/chgnet/model/dynamics.py
@@ -89,7 +89,10 @@ def __init__(
         self.device = device
 
         # Move the model to the specified device
-        self.model = (model or CHGNet.load(verbose=False)).to(self.device)
+        if model is None:
+            self.model = CHGNet.load(verbose=False, use_device=self.device)
+        else:
+            self.model = model.to(self.device)
         self.model.graph_converter.set_isolated_atom_response(on_isolated_atoms)
         self.stress_weight = stress_weight
         print(f"CHGNet will run on {self.device}")