update jax and numpy backend

anyangml · anyangml · commit 2cf38eadbd28 · 2025-12-02T09:30:02.000Z
diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py
@@ -137,7 +137,7 @@ def model_type(self) -> type["DeepEvalWrapper"]:
             return DeepDOS
         elif "dipole" in model_output_type:
             return DeepDipole
-        elif "polar" in model_output_type:
+        elif "polar" in model_output_type or "polarizability" in model_output_type:
             return DeepPolar
         elif "wfc" in model_output_type:
             return DeepWFC
diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py
@@ -158,7 +158,7 @@ def model_type(self) -> type["DeepEvalWrapper"]:
             return DeepDOS
         elif "dipole" in model_output_type:
             return DeepDipole
-        elif "polar" in model_output_type:
+        elif "polar" in model_output_type or "polarizability" in model_output_type:
             return DeepPolar
         elif "wfc" in model_output_type:
             return DeepWFC
