specify map_location in mace_create_lammps_model

ACEsuit · Sep 18, 2024 · a2d3dd7 · a2d3dd7
1 parent 158b1f2
commit a2d3dd7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mace/cli/create_lammps_model.py b/mace/cli/create_lammps_model.py
@@ -54,7 +54,10 @@ def select_head(model):
 def main():
     args = parse_args()
     model_path = args.model_path  # takes model name as command-line input
-    model = torch.load(model_path)
+    model = torch.load(
+        model_path,
+        map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    )
     model = model.double().to("cpu")
 
     if args.head is None: