use faster norm check

iProzd · Nov 27, 2024 · 107dbe7 · 107dbe7
1 parent 6b5f959
commit 107dbe7
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -690,11 +690,8 @@ def step(_step_id, task_key="Default") -> None:
                 loss.backward()
                 if self.gradient_max_norm > 0.0:
                     grad_norm = torch.nn.utils.clip_grad_norm_(
-                        self.wrapper.parameters(), self.gradient_max_norm
+                        self.wrapper.parameters(), self.gradient_max_norm, error_if_nonfinite=True
                     )
-                    if not torch.isfinite(grad_norm).all():
-                        # check local gradnorm single GPU case, trigger NanDetector
-                        raise FloatingPointError("gradients are Nan/Inf")
                 with torch.device("cpu"):
                     self.optimizer.step()
                 self.scheduler.step()