Hi,
Thanks for the reply.
I’d like to compute the gradient of each layer, so basically I do a loop over the layers and call autograd. (the layers in the code below are stored in a list)
  var_x = torch.autograd.Variable(x, requires_grad=True)  
  
  for l in range(n):  
    mx = layers[l](var_x)
    
    gradients = torch.autograd.grad(outputs=mx, inputs=var_x,
                                      grad_outputs=torch.ones(mx.size()).to(device), 
                            create_graph=True, retain_graph=True, only_inputs=True)[0]
      
    norm_gradient_velocity[l] = torch.max( torch.norm( gradients.view(gradients.shape[0],-1), p=np.inf, dim=1 ) )


Because it’s long during training to use such loop, I was wondering if it could be done using hooks.