Pytorch、随机数生成器和设备

Question

我总是在我的 Pytorch 笔记本上放一个像这样的单元格：

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
torch.set_default_device(device)

通过这种便捷的方式，我可以使用 GPU（如果系统有 GPU）、Mac 上的 MPS 或普通系统上的 cpu。

现在，我正在尝试使用 Pytorch 生成器：

g = torch.Generator(device=device).manual_seed(1)

然后：

A = torch.randn((3, 2), generator=g)

在我的 Macbook（设备为 MPS）或仅具有 cpu 的系统上没有任何问题。但在我启用 Cuda 的桌面上，我得到：

RuntimeError: Expected a 'cpu' device type for generator but found 'cuda'

有什么解决办法吗？如果我不为生成器指定设备，它将使用 cpu，但是 然后张量

A

也会在 CPU 上创建...

Answer 1

定义张量“A”时，默认在CPU上创建张量。这会导致错误，因为您传入的生成器位于 CUDA 设备上。您可以通过将“device”参数传入 torch.randn() 来解决此问题，如下所示：

device = ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
g = torch.Generator(device=device).manual_seed(1)
A = torch.randn((3, 2), device=device, generator=g)