I am working on implementing this as well. At what point during the training should you check for the gradient?
Currently, I am checking at the end of each epoch by iterating through my models parameters and calling the variable .grad As shown in code below. However, for some reason when I visualize it in Tensorboard all my layers have zero gradients, even though the histograms show that the weights and bias are changing.
loss.backward()
optimizer.step()
optimizer.zero_grad()
for tag, parm in model.named_parameters:
     writer.add_histogram(tag, parm.grad.data.cpu().numpy(), epoch)