No, it is done by adaptive optimizers. Actually, switching to one may solve your problem by itself, as they smooth steps strongly. RMSProp may be the least intrusive one.



 PhysicsIsFun:

And isn’t Tensorboard part of Tensorflow?


Yes, but pytorch supports writing its trace files. I re-checked my older code that produced gradient plots, and it is a bit non-trivial - I used TensorboardLogger and events from ignite library, and Tensor.register_hook callbacks to record gradients for exporting.