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