PyTorch中如何一次性设置默认计算设备为CUDA？

有几种实用方法可以避免每次手动指定CUDA设备：

1. 设置全局默认张量类型 #

直接修改PyTorch的默认张量类型为CUDA版本，后续所有新创建的张量都会自动分配到CUDA设备：

import torch

# 先确认CUDA可用
if torch.cuda.is_available():
    # 设置默认浮点张量为CUDA版本
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
    # 如果需要其他类型，比如双精度张量
    # torch.set_default_tensor_type('torch.cuda.DoubleTensor')

这样不管是用torch.tensor()、torch.rand()还是torch.zeros()创建的张量，都会直接在CUDA上。同时，新定义的神经网络模型参数也会自动加载到CUDA设备，无需额外调用.to('cuda')。

2. 自定义张量创建工具函数 #

如果不想全局修改默认张量类型（避免影响某些不需要在CUDA运行的张量），可以封装一个简单的函数来快速创建CUDA张量：

import torch

def cuda_tensor(data, dtype=None, device=None, requires_grad=False):
    # 优先使用传入的device，默认用cuda（不可用时自动 fallback 到cpu）
    device = device or 'cuda' if torch.cuda.is_available() else 'cpu'
    return torch.tensor(data, dtype=dtype, device=device, requires_grad=requires_grad)

之后创建张量时直接调用cuda_tensor()代替torch.tensor()即可，比如：

x = cuda_tensor([1,2,3])  # 直接在CUDA上创建张量

3. 模型迁移后匹配设备（辅助技巧） #

当你把模型迁移到CUDA后，可以利用模型的设备属性来统一指定张量的创建设备，避免硬编码：

model = MyModel().to('cuda')
# 创建张量时直接复用模型的设备
x = torch.tensor([1,2,3], device=model.device)

这种方式适合需要灵活切换设备的场景，避免全局设置带来的约束。

注意事项 #

• 所有操作前务必用torch.cuda.is_available()确认CUDA设备可用，否则会抛出运行时错误。
• 全局设置默认张量类型后，临时张量也会占用CUDA显存，注意监控显存使用情况，避免溢出。

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