Gradient Flow Summary
❓ Questions and Help
Before asking:

search the issues.   done
search the docs.    done


What is your question?
I want to add a summary where I can track the gradient flow in my model.
The out-of-the-box gradient tracking is not sufficient for me because I need a more customized behavior.
I wanted to do it using a callback, but I couldn't find a way to access the gradients from a callback.
I noticed there is a model hook for model.on_before_zero_grad(optimizer), which seem to have access for the gradients before resting them, but it seem to don't fit my use case, I would want access to training state information as well when logging..
What is the recommended way to go about it?
thanks,
Alex
Code
What have you tried?
I was aiming fo something of this sort:
class TBGradFlowImageSummary(pl.Callback):
    def on_train_batch_end_before_zero_grad(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
        isLastBatch = batch_idx == len(trainer.train_dataloaders) -1
        if isLastBatch:
                pl_module.logger.experiment.log_something()
What's your environment?

OS: [Linux]
Version [0.9.0]
thanks