pytorch - pytorch-lightning 中的正态分布采样

Question

在Pytorch-Lightning中，您通常不必指定 cuda 或 gpu。但是当我想创建一个高斯采样张量时，torch.normal我得到

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

那么，我该如何更改才能torch.normal使 pytorch-lightning 正常工作？因为我在 cpu和gpu上的不同机器上使用代码

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

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

Answer 1

推荐的方法是lights = torch.normal(0, 1, size=[100, 3], device=self.device)在 Lightning 类中执行此操作。你也可以这样做：lights = torch.normal(0, 1, size=[100, 3]).type_as(tensor)，tensorcuda 上的一些张量在哪里。