support dict

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 12, 2024 · b7d2b32 · b7d2b32
1 parent d5e8b49
commit b7d2b32
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/deepmd/dpmodel/common.py b/deepmd/dpmodel/common.py
@@ -172,6 +172,11 @@ def wrapper(self, *args, **kwargs):
             return tuple(
                 safe_cast_array(vv, self.precision, "global") for vv in returned_tensor
             )
+        elif isinstance(returned_tensor, dict):
+            return {
+                kk: safe_cast_array(vv, self.precision, "global")
+                for kk, vv in returned_tensor.items()
+            }
         else:
             return safe_cast_array(returned_tensor, self.precision, "global")