Skip to content

Commit

Permalink
fix(pt): fix clearing the list in set_eval_descriptor_hook (#4534)
Browse files Browse the repository at this point in the history
Fix #4533.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
  - Improved list clearing mechanism in `DPAtomicModel` class
  - Enhanced test coverage for descriptor evaluation in `TestDeepPot`

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
njzjz authored Jan 7, 2025
1 parent 8d4c27b commit c544c94
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 1 deletion.
3 changes: 2 additions & 1 deletion deepmd/pt/model/atomic_model/dp_atomic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,8 @@ def __init__(
def set_eval_descriptor_hook(self, enable: bool) -> None:
"""Set the hook for evaluating descriptor and clear the cache for descriptor list."""
self.enable_eval_descriptor_hook = enable
self.eval_descriptor_list = []
# = [] does not work; See #4533
self.eval_descriptor_list.clear()

def eval_descriptor(self) -> torch.Tensor:
"""Evaluate the descriptor."""
Expand Down
4 changes: 4 additions & 0 deletions source/tests/infer/test_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,10 @@ def test_descriptor(self) -> None:
descpt = self.dp.eval_descriptor(result.coord, result.box, result.atype)
expected_descpt = result.descriptor
np.testing.assert_almost_equal(descpt.ravel(), expected_descpt.ravel())
# See #4533
descpt = self.dp.eval_descriptor(result.coord, result.box, result.atype)
expected_descpt = result.descriptor
np.testing.assert_almost_equal(descpt.ravel(), expected_descpt.ravel())

def test_2frame_atm(self) -> None:
for ii, result in enumerate(self.case.results):
Expand Down

0 comments on commit c544c94

Please sign in to comment.