diff --git a/pytorch_lightning/core/grads.py b/pytorch_lightning/core/grads.py index 7bdc8572b97fb..d98a6efeb85ca 100644 --- a/pytorch_lightning/core/grads.py +++ b/pytorch_lightning/core/grads.py @@ -10,7 +10,7 @@ class GradInformation(nn.Module): def grad_norm(self, norm_type): results = {} total_norm = 0 - for i, p in enumerate(self.parameters()): + for name, p in self.named_parameters(): if p.requires_grad: try: param_norm = p.grad.data.norm(norm_type) @@ -18,7 +18,7 @@ def grad_norm(self, norm_type): norm = param_norm ** (1 / norm_type) grad = round(norm.data.cpu().numpy().flatten()[0], 3) - results['grad_{}_norm_{}'.format(norm_type, i)] = grad + results['grad_{}_norm_{}'.format(norm_type, name)] = grad except Exception: # this param had no grad pass