Does pytorch support calculate loss on model A and calculate grad on Model B

I want to achieve the function as:
output = modelA(input)
loss = loss_function(output, label)
calculate gradient and update params on ModelB

modelA and modelB have the same architecture.
Can I solve this? and how.

Yes, you could manually copy the .grad attributes from modelA to modelB and update modelB's parameters with its optimizer.

1 Like

Thanks for your advice, I have tried that. Specifically, I wonder if there is a method that the loss is calculated on modelA , and the BP of loss can be calculated on modelB.

I don’t think it’s possible to “redirect” the computation graph to another model, so my guess would be to use the approach of copying the gradients.

Thanks for your advice. If convenient, may you help me to reproduce the result in another my topics and solve it. Use the same data, but get different loss