I have a backward hook function with ‘newLayer.register_backward_hook(hook_function)’ where I do not know how to control the inputs to it.  The function contains a line like.
def hook_function(self, grad_input, grad_output):
    self.average = self.average * 0.99 + grad_output[0].sum((0,2,3)) * 0.01

This results in:
RuntimeError: expected device cuda:0 but got device cuda:1
(Pdb) self.average
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], device='cuda:0',
       dtype=torch.float64)
(Pdb) grad_output[0].sum((0,2,3))
tensor([ 0.0508,  0.0492,  0.0512,  0.0487,  0.0517,  0.0483,  0.0522,  0.0479,
        -0.1974, -0.2026], device='cuda:1', dtype=torch.float64)


I need to know either how to define my self.average better or if i am doing something wrong somewhere else and it is grad_output that is wrong.
the above is with
self.average = torch.Tensor(out_channels).zero_().to(gf.device).double()

I have also tried
self.average = nn.Parameter(torch.Tensor(out_channels).zero_().to(gf.device).double())

which results in
TypeError: cannot assign 'torch.cuda.DoubleTensor' as parameter 'normalPassAverageD' (torch.nn.Parameter or None expected)

at the same location