workaround for bfloat16

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Nov 13, 2024 · e43134d · e43134d
1 parent 4ba6e97
commit e43134d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py
@@ -248,6 +248,10 @@ def call(self, x: np.ndarray) -> np.ndarray:
             if self.b is not None
             else xp.matmul(x, self.w)
         )
+        if y.dtype != x.dtype:
+            # workaround for bfloat16
+            # https://github.com/jax-ml/ml_dtypes/issues/235
+            y = xp.astype(y, x.dtype)
         y = fn(y)
         if self.idt is not None:
             y *= self.idt