Wait... something about your code has just broken my tiny brain! Why is there no call to `optimizer.step()` or `loss.backward()`... ?? Also, it looks like your using lucidrains' EMA lib, but am I correct in assuming that?