pgujjar:

I suspect this is because higher takes a snapshot of model B only


with higher.innerloop_ctx(student, student_optim, copy_initial_weights=False) as (fmodel, diffopt):
   pseudo_labels = teacher(x)

   logits = fmodel(x)
   loss1 = criterion(logits, pseudo_labels)
   diffopt.step(loss1)

   logits = fmodel(x)
   loss2 = criterion(logits, y_val)
   grad = torch.autograd.grad(loss2, teacher.parameters())

I agree. It seems to be working when replacing modelB with fmodel  in second the update.