For some reason I need to call loss.backward() to tarin a special parameter in the middle layer of my model, but for the sake of saving calculation expenses I don’t want the backpropagation to continue once I get the gradient for my parameter, otherwise I’ll run out of my kaggle-GPU free usage . So is there any way to manually stop backpropagation after It was done for a specific parameter or a module?
Hi,
You can use autograd.grad(loss, input)
and it will return gradients for that input and only perform the necessary computations to get that value.
If you prefer for the .grad
fields to be populated, you can use loss.backward(inputs=input)
.
1 Like
Helped me out!Thanks alt