Removed redundant computations in clip_gradients that slowed down the gradient clipping. (#1523)

Fixes #1522
2026-09-09 11:32:07 +08:00 · 2020-04-18 23:07:15 -04:00
parent a22a8142ac
commit e02146943d
1 changed files with 1 additions and 1 deletions
@@ -40,7 +40,7 @@ class TrainerTrainingTricksMixin(ABC):
                device = parameters[0].device
                total_norm = torch.zeros([], device=device if parameters else None)
                for p in parameters:
-                    param_norm = p.grad.data.norm(norm_type) ** norm_type
+                    param_norm = p.grad.data.pow(norm_type).sum()
                    total_norm.add_(param_norm)
                total_norm = (total_norm ** (1. / norm_type))
            eps = EPSILON_FP16 if self.precision == 16 else EPSILON