Thank you. But,
Is this the same as setting modelB to eval mode?
and,
If I want to update the weights of both the models, can I achieve this by creating two optimizers?
optimizer_A = torch.optim.Adam(modelA.parameters())
optimizer_B = torch.optim.Adam(modelB.parameters())
And in training loop
optimizer_A.zero_grad()
optimizer_B.zero_grad()
opA = modelA(inp)
opB = modelB(opA)
loss = criterion(opB, GT)
loss.backward()
optimizer_B.step()
optimizer_A.step()