Skip to content

Commit 2d5fa3c

Browse files
authored
Fix: model_output_type unify name (#5069)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved model type detection to also recognize "polarizability" as equivalent to "polar", ensuring models that report polarizability are correctly handled during inference and expanding compatibility with additional model variants. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 4324445 commit 2d5fa3c

File tree

3 files changed

+3
-3
lines changed

3 files changed

+3
-3
lines changed

deepmd/dpmodel/infer/deep_eval.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def model_type(self) -> type["DeepEvalWrapper"]:
137137
return DeepDOS
138138
elif "dipole" in model_output_type:
139139
return DeepDipole
140-
elif "polar" in model_output_type:
140+
elif "polar" in model_output_type or "polarizability" in model_output_type:
141141
return DeepPolar
142142
elif "wfc" in model_output_type:
143143
return DeepWFC

deepmd/jax/infer/deep_eval.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ def model_type(self) -> type["DeepEvalWrapper"]:
158158
return DeepDOS
159159
elif "dipole" in model_output_type:
160160
return DeepDipole
161-
elif "polar" in model_output_type:
161+
elif "polar" in model_output_type or "polarizability" in model_output_type:
162162
return DeepPolar
163163
elif "wfc" in model_output_type:
164164
return DeepWFC

deepmd/pt/infer/deep_eval.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ def model_type(self) -> type["DeepEvalWrapper"]:
248248
return DeepDOS
249249
elif "dipole" in model_output_type:
250250
return DeepDipole
251-
elif "polar" in model_output_type:
251+
elif "polar" in model_output_type or "polarizability" in model_output_type:
252252
return DeepPolar
253253
elif "global_polar" in model_output_type:
254254
return DeepGlobalPolar

0 commit comments

Comments
 (0)