perf: optimize training loop

iProzd · Nov 27, 2024 · 6b5f959 · 6b5f959
1 parent 9840ae6
commit 6b5f959
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -662,7 +662,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: