Skip to content

Commit

Permalink
perf: optimize training loop
Browse files Browse the repository at this point in the history
  • Loading branch information
caic99 authored Nov 26, 2024
1 parent e7ad8dc commit 4823dd3
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion deepmd/pt/train/training.py
Original file line number Diff line number Diff line change
Expand Up @@ -656,7 +656,8 @@ def step(_step_id, task_key="Default") -> None:
# PyTorch Profiler
if self.enable_profiler or self.profiling:
prof.step()
self.wrapper.train()
if not self.wrapper.training:
self.wrapper.train()
if isinstance(self.lr_exp, dict):
_lr = self.lr_exp[task_key]
else:
Expand Down

0 comments on commit 4823dd3

Please sign in to comment.