Solved by ptrblck in post #2 
              
              
                I don’t know how self.average is initialized, but would assume this should work: 
self.average = self.average.to(grad_output[0].device) * 0.99 + grad_output[0].sum((0,2,3)) * 0.01

Could you check it and see, if you are still getting an error? 
In that case, could you post a small code snippet to re…