JuJu:

@staticmethod def backward(ctx, grad_output): return (ctx.input * ctx.yt).backward() / ctx.running_ema, None, None


I think, you do not need .backward() call inside backward method.
    @staticmethod
    def backward(ctx, grad_output):
        return (ctx.input * ctx.yt) / ctx.running_ema, None, None

I did not read the paper. But this will resolve the error atleast