Hi Guys,
I keep getting this error with a custom module I’m trying to create:
"expected device cuda:0 and dtype Float but got device cpu and dtype Float"
From what I can understand, I need to use register_buffer on the offending tensor so that it gets set to CUDA when I use model.CUDA(). I think I’m doing that below, but I keep getting the same error.
class NoiseInjection(nn.Module):
def __init__(self, shape):
super(NoiseInjection, self).__init__()
self.injected_noise = self.build_noise_injection(shape)
def build_noise_injection(self, shape):
noise = torch.randn([shape[0], 1, shape[2] , shape[3]], requires_grad=False) #NCHW
self.register_buffer('noise', noise)
weights = nn.Parameter(torch.zeros([1, shape[1], 1 , 1])) # different weight for each channel
noise_weighted = weights * self.noise
return noise_weighted
def forward(self,x):
out = x + self.injected_noise
return out
Anyone know what’s going on?