From 5bef046759c0ad202ed4b77803bad548671fce10 Mon Sep 17 00:00:00 2001
From: Rohit Gupta <rohitgr1998@gmail.com>
Date: Wed, 13 Jan 2021 05:06:13 +0530
Subject: [PATCH] [Docs] fix on_after_backward example (#5278)

* fix on_after_backward docs

* doc fix
---
 pytorch_lightning/core/hooks.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py
index f441033487..e8d7699cd1 100644
--- a/pytorch_lightning/core/hooks.py
+++ b/pytorch_lightning/core/hooks.py
@@ -307,10 +307,9 @@ class ModelHooks:
                 if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
                     params = self.state_dict()
                     for k, v in params.items():
-                        grads = v
-                        name = k
-                        self.logger.experiment.add_histogram(tag=name, values=grads,
-                                                             global_step=self.trainer.global_step)
+                        self.logger.experiment.add_histogram(
+                            tag=k, values=v.grad, global_step=self.trainer.global_step
+                        )
 
         """