Logging gradients and a possible bug in the docs?
TL;DR Question:  How do we correctly log gradients to tensorboard?
I am trying to log gradients to tensorboard to track NaN loss in a speech application. The docs suggest using on_after_backward along with this code that appears to be incorrect:
# 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)

This code seems to log the weights instead of gradients (assuming lightning state_dict is the same structure as pytorch). I'm happy to fix it and submit a PR as long as I'm not mistaken. I would log the weights like this...
# 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 name, weight in params.items():
         self.logger.experiment.add_histogram(tag=name, values=weight,
                                                             global_step=self.trainer.global_step)

but for gradients, if I try to log weights.grad, I get None. How do we correctly log gradients to tensorboard?
# 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 name, weight in params.items():
           self.logger.experiment.add_histogram(tag=name, values=weight.grad,
                                                             global_step=self.trainer.global_step)