avoid import torch for type check in MontyEncoder

materialsvirtuallab · Oct 20, 2024 · a7752ee · a7752ee
1 parent 1270c7b
commit a7752ee
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/monty/json.py b/src/monty/json.py
@@ -586,8 +586,8 @@ def default(self, o) -> dict:
         if isinstance(o, Path):
             return {"@module": "pathlib", "@class": "Path", "string": str(o)}
 
-        if torch is not None and isinstance(o, torch.Tensor):
-            # Support for Pytorch Tensors.
+        # Support for Pytorch Tensors
+        if _check_type(o, "torch.Tensor"):
             d: dict[str, Any] = {
                 "@module": "torch",
                 "@class": "Tensor",