I have two models A and B. The output of model A is fed to model B and the loss is calculated(example below)
opA = modelA(inp)
opB = modelB(opA)
loss = criterion(opB, GT)
I want to update weights of model A only. How can I achieve this?