Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions deepmd/pt_expt/infer/deep_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -366,10 +366,15 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None:
# eager inference). Drop the latter and unwrap the former.
cleaned: dict[str, Any] = {}
compiled_marker = ".compiled_forward_lower."
# Per-task buffer copies registered on _CompiledModel (bias_atom_e,
# case_embd) — real values live on the original model's fitting net.
task_buf_marker = "._task_"
wrapper_infix = ".original_model."
for key, value in state_dict.items():
if compiled_marker in key:
continue
if task_buf_marker in key:
continue
if wrapper_infix in key:
key = key.replace(wrapper_infix, ".", 1)
cleaned[key] = value
Expand Down
Loading
Loading