Using output of one network in another with only one loss function

autograd
1

I am using two neural networks say NN1 and NN2. Output of NN1 is used in forward function of NN2 and output of forward function of NN2 is the loss required. So in the training loop if I write

loss = NN2(data) 
loss.backward()

Will this update the weights of both NN1 and NN2?

2

If I’m understanding your use case correctly, you are sequentially calling NN1 and NN2.
If that’s the case, your code should be fine and gradients should be calculated for both models.
Basically you can treat models (like NN1 and NN2) as layers, since both are derived from nn.Module.
There will be no difference between sequentially calling e.g. nn.Linear layer instances and your models.

1 Like