如何解决TorchScript模型CUDA设备不一致的问题

在PyTorch的模型部署场景中，TorchScript作为将PyTorch模型序列化、便于跨环境部署的工具，被广泛应用。但很多开发者在使用TorchScript模型做CUDA推理时，会遇到模型设备与输入数据设备不匹配的报错，导致推理无法执行。

问题常见表现

当TorchScript模型的权重存放在CUDA设备0，而输入数据被放到CUDA设备1，或者模型在CPU上、输入在CUDA上时，运行推理通常会抛出类似以下的错误：

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

还有可能出现模型加载后默认在CPU，但是输入数据已经转到CUDA，执行推理时也会出现设备不匹配的提示。

问题产生原因

出现CUDA设备不一致的核心原因是模型权重和输入数据的存放设备没有统一，常见场景有以下几种：

• 加载TorchScript模型时没有指定目标设备，模型默认加载到保存时的原始设备，和当前输入数据的设备不同
• 多GPU环境下，模型被加载到某个GPU，但是输入数据被错误地分配到了其他GPU
• 模型保存时是在CPU上，加载后没有转移到CUDA，而输入数据已经提前转到了CUDA设备

解决方法

方法一：加载模型时指定目标设备

在加载TorchScript模型的时候，可以通过map_location参数直接指定模型加载到目标CUDA设备，从源头保证模型权重的设备和预期一致。

import torch

# 指定模型加载到cuda:0设备
model = torch.jit.load("your_torchscript_model.pt", map_location="cuda:0")

# 构造输入数据，同样放到cuda:0
input_tensor = torch.randn(1, 3, 224, 224).to("cuda:0")

# 执行推理
output = model(input_tensor)
print(output.shape)

方法二：加载后手动转移模型到目标设备

如果已经加载了模型，也可以通过to方法将整个模型转移到目标CUDA设备，注意该方法会将模型所有参数都转移到指定设备。

import torch

# 先加载模型，默认可能加载到CPU或者其他设备
model = torch.jit.load("your_torchscript_model.pt")

# 将模型转移到cuda:1设备
model = model.to("cuda:1")

# 输入数据也需要同步到cuda:1
input_tensor = torch.randn(1, 3, 224, 224).to("cuda:1")

output = model(input_tensor)

方法三：统一输入数据和模型的设备

如果模型已经固定了设备，也可以将输入数据转移到模型所在的设备，通过模型的device属性可以获取当前模型所在的设备。

import torch

# 加载模型，假设加载到了cuda:0
model = torch.jit.load("your_torchscript_model.pt", map_location="cuda:0")

# 获取模型所在的设备
model_device = next(model.parameters()).device

# 将输入数据转移到模型所在设备
input_tensor = torch.randn(1, 3, 224, 224).to(model_device)

output = model(input_tensor)

多GPU场景的注意事项

如果是在多GPU环境下使用TorchScript模型，需要提前确认当前可用的GPU编号，避免指定不存在的CUDA设备。可以通过torch.cuda.device_count()查看可用GPU数量，通过torch.cuda.current_device()查看当前默认GPU。

另外，如果模型是通过torch.nn.DataParallel训练后转成的TorchScript，需要注意DataParallel包装的模型序列化后可能会带有module.前缀，加载后可能需要额外处理，不过这种情况较少出现在纯TorchScript部署场景中。

验证设备是否一致

在推理前可以加一段验证代码，确认模型和输入的设备是否统一，避免运行时报错：

import torch

model = torch.jit.load("your_torchscript_model.pt", map_location="cuda:0")
input_tensor = torch.randn(1, 3, 224, 224).to("cuda:0")

model_device = next(model.parameters()).device
input_device = input_tensor.device

if model_device == input_device:
    print("模型和输入设备一致，可以执行推理")
    output = model(input_tensor)
else:
    print(f"设备不一致，模型在{model_device}，输入在{input_device}")

通过上述方法，基本可以覆盖所有TorchScript模型CUDA设备不一致的场景，开发者可以根据实际的部署环境选择对应的解决方式。

TorchScriptCUDAPyTorch模型推理修改时间：2026-06-25 17:48:36