How to parallelize this training?

Hello, I love your method but am trying to find a more efficient way to run it. In fact, I am willing to contribute to the project to make training with this algorithm more realistic. Please let me know if you have already put thought into making the distillation algorithm more efficient and what ideas you have already explored!