如何判断PyTorch模型当前运行在cpu还是cuda设备上？

核心前提 #

nn.Module本身没有内置的设备标识属性（比如类似张量的is_cuda或直接返回device的属性），因为PyTorch将nn.Module定义为逻辑容器，而非持有实际数据的张量实体，它的设备状态由其管理的可训练参数、非训练缓冲区（如BatchNorm的running_mean）的设备共同决定。

常用判断方法 #

1. 常规场景：取第一个参数的设备（最便捷） #

绝大多数使用场景下，不会刻意将同一个模型的参数拆分到不同设备，因此只需取第一个可训练参数的设备即可等价为整个模型的设备：

# 直接获取模型设备
model_device = next(model.parameters()).device

# 后续初始化其他张量时直接指定设备，无需手动区分CPU/CUDA构造函数
# 原来需要判断设备后分别调用torch.Tensor / torch.cuda.Tensor的写法，现在简化为：
new_tensor = torch.tensor(your_data, device=model_device)

注意：如果是未添加任何参数的空nn.Module（仅定义了__init__未定义层），model.parameters()会返回空迭代器，此时会触发报错，这类空模型默认运行在CPU设备上，可自行加异常处理兜底。

2. 严谨场景：校验所有参数/缓冲区的设备一致性 #

如果需要确认模型没有被手动修改过参数设备（排除模型并行/分布式的特殊场景），可以遍历所有参数和缓冲区，统一校验设备是否一致：

def get_consistent_model_device(model: nn.Module):
    # 收集所有参数、缓冲区的设备
    device_set = {param.device for param in model.parameters()}
    device_set.update({buf.device for buf in model.buffers()})
    
    if len(device_set) == 1:
        return next(iter(device_set))
    else:
        raise RuntimeError(f"模型参数分布在多个设备上：{device_set}，请检查是否有手动调整设备的操作")

补充说明 #

不需要手动通过is_cuda属性判断是CPU还是CUDA设备，直接复用获取到的device对象给其他张量指定device参数，PyTorch会自动生成对应设备的张量，无需区分torch.Tensor（CPU默认）或torch.cuda.Tensor的构造逻辑。

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