fix: fix for __array__ removal from array-api-strict

See: https://github.com/data-apis/array-api-strict/blob/main/docs/changelog.md#major-changes Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 7, 2024 · cde3039 · cde3039
1 parent 3236db5
commit cde3039
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/deepmd/dpmodel/common.py b/deepmd/dpmodel/common.py
@@ -105,6 +105,10 @@ def to_numpy_array(x: Any) -> Optional[np.ndarray]:
     """
     if x is None:
         return None
+    if x.device == "cpu":
+        # dlpack needs the device to be the same
+        return np.from_dlpack(x)
+    # asarray is not within Array API standard, so may fail
     return np.asarray(x)