火炬闪电的正态分布采样

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，使火把-闪电正常工作呢？因为我在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)。你也可以这样做：lights = torch.normal(0, 1, size=[100, 3]).type_as(tensor)，其中tensor是一些张量，在库达上。