I run your code,meet the error as flows:
square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)
RuntimeError: cuda runtime error (59) : device-side assert triggered at /pytorch/aten/src/THC/generic/THCTensorMathPairwise.cu:81
could you help me?thanks!