PyTorch生成器指定CUDA设备报错，如何保持张量在目标设备生成？

问题分析 #

你遇到的RuntimeError: Expected a 'cpu' device type for generator but found 'cuda'报错，根源在于PyTorch中CUDA设备的随机数生成器实现逻辑与CPU、MPS存在差异：CPU和MPS支持直接通过torch.Generator(device=device)指定设备创建生成器，但CUDA生成器需要关联CUDA流，直接用通用方式创建会触发设备不匹配错误。

解决方案 #

方案一：根据设备类型适配生成器创建逻辑 #

针对不同设备类型选择对应的生成器创建方式，既保证生成器与设备匹配，又能让张量默认创建在目标设备上：

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

# 根据设备类型创建生成器
if device == "cuda":
    g = torch.cuda.Generator().manual_seed(1)
else:
    g = torch.Generator(device=device).manual_seed(1)

# 生成张量，自动使用默认设备
A = torch.randn((3, 2), generator=g)

• 对于CUDA设备，使用torch.cuda.Generator()创建专属生成器，它会自动关联当前CUDA设备的默认流，彻底避免设备不匹配问题；
• 对于CPU和MPS设备，依然保留原有的生成器创建方式，保持跨设备兼容性。

方案二：显式指定张量设备（兼容通用生成器） #

如果想统一使用torch.Generator()的创建方式，可以在生成张量时显式指定设备，即使生成器在CPU上，张量也会直接创建在目标设备：

g = torch.Generator(device="cpu").manual_seed(1)
A = torch.randn((3, 2), generator=g, device=device)

注意：这种方式会先在CPU生成随机数再转移到GPU，性能略低于方案一，适合对性能要求不高的场景。

验证说明 #

• 方案一中，CUDA生成器与默认设备完全匹配，torch.randn会直接在CUDA设备上生成张量，无额外数据转移开销；
• 两种方案都能保证在CPU、MPS、CUDA设备上均正常运行，且张量默认（或显式）创建在目标设备。

内容的提问来源于stack exchange，提问作者MadHatter