fix grad norm formula

Author: ivannz

pytorch_lightning/core/grads.py

@@ -3,28 +3,25 @@
 """
 from typing import Dict
 
-from torch import nn
+import torch
 
 
-class GradInformation(nn.Module):
+class GradInformation(torch.nn.Module):
 
     def grad_norm(self, norm_type: float) -> Dict[str, int]:
-        results = {}
-        total_norm = 0
+        norms, all_norms = {}, []
         for name, p in self.named_parameters():
-            if p.requires_grad:
-                try:
-                    param_norm = p.grad.data.norm(norm_type)
-                    total_norm += param_norm ** norm_type
-                    norm = param_norm ** (1 / norm_type)
-
-                    grad = round(norm.data.cpu().numpy().flatten()[0], 3)
-                    results['grad_{}_norm_{}'.format(norm_type, name)] = grad
-                except Exception:
-                    # this param had no grad
-                    pass
-
-        total_norm = total_norm ** (1. / norm_type)
-        grad = round(total_norm.data.cpu().numpy().flatten()[0], 3)
-        results['grad_{}_norm_total'.format(norm_type)] = grad
-        return results
+            if p.grad is None:
+                continue
+
+            param_norm = float(p.grad.data.norm(norm_type))
+            norms[f'grad_{norm_type}_norm_{name}'] = round(param_norm, 3)
+
+            all_norms.append(param_norm)
+
+        total_norm = 0.
+        if all_norms:
+            total_norm = float(torch.tensor(all_norms).norm(norm_type))
+        norms[f'grad_{norm_type}_norm_total'] = round(total_norm, 3)
+
+        return norms