update example (#5753)

pytorch_lightning/core/hooks.py

@@ -305,8 +305,7 @@ class ModelHooks:
             def on_after_backward(self):
                 # example to inspect gradient information in tensorboard
                 if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
-                    params = self.state_dict()
-                    for k, v in params.items():
+                    for k, v in self.named_parameters():
                         self.logger.experiment.add_histogram(
                             tag=k, values=v.grad, global_step=self.trainer.global_step
                         )