I am trying to implement the Meta-Pseudo Label paper (https://arxiv.org/pdf/2003.10580v2.pdf 2)  where a teacher model is updated based on a student model’s performance on a validation set. It would look something like this:
modelA_optim = optim.Adam(modelA.parameters(), lr=0.01)
modelB_optim = optim.Adam(modelB.parameters(), lr=0.01)

pseudo_labels = modelA(X_train)

# Train model B on pseudo labels
logits = modelB(X_train)
loss1 = CrossEntropy(logits, pseudo_labels)
loss1.backward()
modelB_optim.step()

# Update model A with model B's loss on a validation set
logits = modelB(X_val)
loss2 = CrossEntropy(logits, y_val)
loss2.backward()
modelA_optim.step()

This in iteself wouldn’t update model A’s weights as the gradient cannot backpropate all the way through model B and eventually to model A through the graient of loss1.
I have tried using the higher library to make model B a functional one, thereby being able to backpropate through its one updat (becuase of loss1, model B’s weights_t will now be weights_t+1). However, torch.autograd.grad(loss, modelA.parameters()) yields zero gradients for model A. I suspect this is because the backward graph is broken at the loss1 point of connection between model A and model B. Is there a way to address this?!