如何在GPU上进行PyTorch分发

Question

是否有可能使PyTorch发行版直接在GPU上创建其样本。

如果我这样做

from torch.distributions import Uniform, Normal
normal = Normal(3, 1)
sample = normal.sample()

然后样本将在CPU上。当然，可以执行sample = sample.to(torch.device("cuda"))在GPU上进行制作。但是有没有一种方法可以使样本直接进入GPU，而无需先在CPU上创建它？

PyTorch分配继承自Object，而不是nn.Module，因此它没有to方法将分配实例放在GPU上。

有什么想法吗？

Answer 1

发行版使用重新参数化技巧。因此，将GPU上的大小为0的张量提供给分布构造函数。如下：

normal = Normal(torch.tensor(0).to(device=torch.device("cuda")), torch.tensor(1).to(device=torch.device("cuda")))