Hey!
I’m also trying to implement meta pseudo labels - this is the training loop I have, but not sure if teacher_loss.backward() is updating the parameters in the right way.
Any advice would be much appreciated!
for i in range(epochs):
    print(i)
    with higher.innerloop_ctx(student, student_optim, copy_initial_weights=False) as (fstudnet, diffopt):
        
        pseudo_labels = teacher(x_train) # labels from teacher

        logits = fstudnet(x_train)
        student_loss = student_loss_func(logits, pseudo_labels)
        
        diffopt.step(student_loss) # zero grad done implicitly 

        logits = fstudnet(x_val)
        teacher_loss = teacher_loss_func(logits, y_val)
        
#         grad = torch.autograd.grad(teacher_loss, teacher.parameters()) 
#         not sure how to pass grad explicitly to model parameters

        teacher_loss.backward()        
        teacher_optim.step() 
        teacher_optim.zero_grad()