Teacher-Student learning Framework: using EMA to distil knowledge from one model to the other

Hello there!
I have two models that are identical and I’m trying to update one of them using the other in an EMA manner.
like this:

I’ve attempted to do this using the following code:

    with torch.no_grad():
        for teacher_param, student_param in zip(teacher.parameters(), student.parameters()):
            teacher_param.data *= tau
            teacher_param.data += (1 - tau) * student_param.data

Is this the correct approach to achieve my goal?
Thank you so much for your help!

I found a paper called Unbiased Teacher for Semi-Supervised Object Detection that implemented this technique. They have publicly available code, and here is a quick url to see how they implemented it:

https://github.com/facebookresearch/unbiased-teacher/blob/226f8c4b7b9c714e52f9bf640da7c75343803c52/ubteacher/engine/trainer.py#L632