Docstring for `on_after_backward` 
📚 Documentation
Hi !
In the docstring for on_after_backward there is a puzzling piece of code that is suggested (link) :
# 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():
                        grads = v
                        name = k
                        self.logger.experiment.add_histogram(tag=name, values=grads,
                                                             global_step=self.trainer.global_step)

It isn't reported in Pytorch documentation that enumerating the state dict key-values gives the gradient: it is usually used to load a saved model weights (thus grads would be the weights and not the grads).
Adding a reference (which I couldn't find) would probably help pick up the logic behind it.