Fix model loading with map_location.

materialsvirtuallab · Aug 6, 2023 · 4d53fe7 · 4d53fe7
1 parent 4961e1c
commit 4d53fe7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 10 deletions.
diff --git a/matgl/graph/compute.py b/matgl/graph/compute.py
@@ -104,11 +104,7 @@ def compute_theta_and_phi(edges: dgl.udf.EdgeBatch):
     triple_bond_lengths (torch.tensor):
     """
     angles = compute_theta(edges, cosine=True)
-    angles.update(
-        {
-            "phi": torch.zeros_like(angles["cos_theta"]),
-        }
-    )
+    angles["phi"] = torch.zeros_like(angles["cos_theta"])
     return angles
 
 

diff --git a/matgl/utils/io.py b/matgl/utils/io.py
@@ -111,11 +111,9 @@ def load(cls, path: str | Path | dict, **kwargs):
 
         _check_ver(cls, model_data)
 
-        if not torch.cuda.is_available():
-            state = torch.load(fpaths["state.pt"], map_location=torch.device("cpu"))
-        else:
-            state = torch.load(fpaths["state.pt"])
-        d = torch.load(fpaths["model.pt"])
+        map_location = torch.device("cpu") if not torch.cuda.is_available() else None
+        state = torch.load(fpaths["state.pt"], map_location=map_location)
+        d = torch.load(fpaths["model.pt"], map_location=map_location)
 
         # Deserialize any args that are IOMixIn subclasses.
         for k, v in d.items():