Solved by ilkarman in post #3 
              
              
                If you use tensorboardX you could also log the grades to tensorboard to visualise: 
from tensorboardX import SummaryWriter
logger = SummaryWriter(LOG_DIR)
    
def log_gradients_in_model(model, logger, step):
    for tag, value in model.named_parameters():
        if value.grad is not None:
        …