deepmodeling / deepmodeling/deepmd-kit

fix(pt-expt): preserve DOS width in exported models

Open
#5,949 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

## Problem

PyTorch-Exportable DOS models do not preserve or report their configured output width after export.

A model configured with `numb_dos: 2` can train and export to `.pte`, but the pt_expt model/evaluator reports `get_numb_dos() == 0`. `DeepDOS(...).eval()` then fails while reshaping the two predicted values into an output width of zero.

The current gaps are:

- `deepmd/pt_expt/model/dos_model.py::DOSModel` does not expose the configured DOS width as the DP-model DOS class does.
- `deepmd/pt_expt/infer/deep_eval.py::get_numb_dos()` returns `0` unconditionally.

## Expected behavior

Exported PyTorch-Exportable DOS models should preserve `numb_dos`, and `DeepDOS`/`dp test` should recover that value and evaluate the exported model successfully.

## Context

Discovered while reviewing #5929. Until this is fixed, the PyTorch-Exportable DOS training/export/testing workflow should not be advertised as supported.

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.