Updating the parameters of a few nodes in a pre-trained network during training

I think it’s better to avoid modifying the gradients in place, but instead return a new gradient, as explained in the docs.
So a modified version would be to pass a hook such as:

def my_hook(grad):
    grad_clone = grad.clone()
    grad_clone[:, 2500:] = 0
    return grad_clone
2 Likes