如何在 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")))