I don’t know how each of the posted methods is implemented and would guess that new tensors are created internally without using the device argument of an already defined and registered parameter or buffer.
bias_zero is properly registered as a buffer and will thus be pushed to the right device.
However, self.bias_sampler.sample() could create a new CPUTensor, which would then create the issue.
In case you’ve implemented these methods, make sure that new tensors are pushed to the right device e.g. via:
device = next(self.paramters()).device
new_tensor = torch.randn(..., device=device)