Normal distribution sampling in pytorch-lightning

Question

In Pytorch-Lightning you usually never have to specify cuda or gpu. But when I want to create a gaussian sampled Tensor using torch.normal I get

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

so, how do I have to change torch.normal such that pytorch-lightning works properly? Since I use the code on different machines on cpu and on gpu

centers = data["centers"] #already on GPU... sometimes...

lights = torch.normal(0, 1, size=[100, 3])
lights += centers

Answer 1

The recommended way is to do lights = torch.normal(0, 1, size=[100, 3], device=self.device) if this is inside lightning class. You could also do: lights = torch.normal(0, 1, size=[100, 3]).type_as(tensor), where tensor is some tensor which is on cuda.