I have a neural network divided into two parts. One part computes the first n layers and the other part computes the rest of the layers.
I can obtain the gradients from a single model this way:
grads = torch.autograd.grad(loss, model.parameters())

The thing is I need the gradients of the two parts separately. The main issue I am facing seems to be continuing the first part’s gradients backwards from the second part.
I would appreciate any help as to how I can achieve this. Thanks