Skip to content

Commit

Permalink
fix(pt): detach computed descriptor tensor to prevent OOM (deepmodeli…
Browse files Browse the repository at this point in the history
…ng#4547)

Fix deepmodeling#4544.

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

## Summary by CodeRabbit

- **Optimization**
- Improved descriptor handling in the atomic model to enhance
computational efficiency and memory management during backpropagation.

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

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
njzjz authored Feb 5, 2025
1 parent 8104841 commit 228062c
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion deepmd/pt/model/atomic_model/dp_atomic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,7 @@ def forward_atomic(
)
assert descriptor is not None
if self.enable_eval_descriptor_hook:
self.eval_descriptor_list.append(descriptor)
self.eval_descriptor_list.append(descriptor.detach())
# energy, force
fit_ret = self.fitting_net(
descriptor,
Expand Down

0 comments on commit 228062c

Please sign in to comment.