在Google Colab训练PyTorch神经ODE时遇RuntimeError求助

数据老陈

2026-7-8

问题描述

在Google Colab中使用PyTorch训练神经ODE网络时，执行optim.step()触发如下错误：

RuntimeError: Tensors of the same index must be on the same device and the same dtype except `step` tensors that can be CPU and float32 notwithstanding

完整回溯信息：

Traceback (most recent call last):
  File "/content/gdrive/MyDrive/damage_neural/neural.py", line 370, in <module>
    optim.step()
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py", line 373, in wrapper
    out = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py", line 76, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/adam.py", line 163, in step
    adam(
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/adam.py", line 311, in adam
    func(params,
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/adam.py", line 474, in _multi_tensor_adam
    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
  File "/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py", line 397, in _group_tensors_by_device_and_dtype
    return _group_tensors_by_device_and_dtype(tensorlistlist, with_indices)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/_foreach_utils.py", line 42, in _group_tensors_by_device_and_dtype
    torch._C._group_tensors_by_device_and_dtype(tensorlistlist, with_indices).items()
RuntimeError: Tensors of the same index must be on the same device and the same dtype except `step` tensors that can be CPU and float32 notwithstanding

错误出现在训练循环的optim.step()行：

for epoch in tqdm(range(n_epochs)):
        optim.zero_grad()
        r = ode.odeint_adjoint(model, y[0],t,block_size=time_chunk_size)
        # loss_val = loss(r[...,0], y[...,0]) + loss(r[...,1], y[...,1]) + loss(r[...,2], y[...,2]) + loss(r[...,3], y[...,3])
        loss_val = loss(r, y)
        loss_hist.append(loss_val.cpu().detach().numpy())
        loss_val.backward()
------> optim.step()
        writer.add_scalar('Loss', loss_val, epoch)
        if loss_val < best_loss:
            best_loss = loss_val
            torch.save(model.state_dict(), 'best_model.pt')

用户已尝试将数据移至同一设备：

d1 = (data_list[1] - torch.min(data_list[1])) / (torch.max(data_list[1]) - torch.min(data_list[1])).cuda()
    d2 = (data_list[3] - torch.min(data_list[3])) / (torch.max(data_list[3]) - torch.min(data_list[3])).cuda()
    d3 = (data_list[5] - torch.min(data_list[5])) / (torch.max(data_list[5]) - torch.min(data_list[5])).cuda()
    d4 = (data_list[6] - torch.min(data_list[6])) / (torch.max(data_list[6]) - torch.min(data_list[6])).cuda()
    d5 = (data_list[7] - torch.min(data_list[7])) / (torch.max(data_list[7]) - torch.min(data_list[7])).cuda()
    d6 = (data_list[16] - torch.min(data_list[16])) / (torch.max(data_list[16]) - torch.min(data_list[16])).cuda()

    # Create the state variable
    y = torch.stack([d2, d4, d3, d6]).cuda() # 4 x 39 x 1000
    y = y.permute(2,1,0).cuda() # 1000 x 39 x 4

    t = t.permute(1,0).cuda()

    d5= d5.permute(1,0).cuda()
    d1 = d1.permute(1,0).cuda()

    # Create the model
    model = NeuralODEModel(w_in, b_in, w_hid, b_hid, w_out, b_out, d5, d1, t).to(device)
    
    # Integrate the ODE
    r = ode.odeint_adjoint(model, y[0],t,block_size=time_chunk_size).cuda() # 1000 x 39 x 4

    r = r.permute(1,0,2).cuda() # 39 x 1000 x 4

该代码在本地Linux机器运行正常，但在Colab中报错。

解决方案

这个错误的核心是优化器中的参数张量存在设备/ dtype不统一的情况，以下是具体修复步骤：

调整优化器初始化时机
必须先将模型移至目标设备（GPU），再初始化优化器。如果先初始化优化器再移动模型，优化器内保存的参数仍为CPU张量，会导致设备不匹配。
修正示例：
```
# 先把模型移到指定设备
model = NeuralODEModel(...).to(device)
# 再初始化优化器
optim = torch.optim.Adam(model.parameters(), lr=1e-3)
```

统一模型传入的固定张量设备与 dtype
模型初始化时传入的d5、d1、t等固定张量，需确保它们的设备和 dtype 与模型完全一致。可以在传入模型前统一处理：

d5 = d5.permute(1,0).to(device, dtype=torch.float32)
d1 = d1.permute(1,0).to(device, dtype=torch.float32)
t = t.permute(1,0).to(device, dtype=torch.float32)
model = NeuralODEModel(w_in, b_in, w_hid, b_hid, w_out, b_out, d5, d1, t).to(device)

也可以在模型内部自动对齐设备：

class NeuralODEModel(nn.Module):
    def __init__(self, w_in, b_in, w_hid, b_hid, w_out, b_out, d5, d1, t):
        super().__init__()
        # 可训练参数定义...
        # 将固定张量移至模型所在设备
        self.d5 = d5.to(self.device)
        self.d1 = d1.to(self.device)
        self.t = t.to(self.device)

禁用多张量优化器（临时 workaround）
部分PyTorch版本的多张量Adam实现对设备一致性检查更严格，可通过设置foreach=False禁用该特性：
```
optim = torch.optim.Adam(model.parameters(), lr=1e-3, foreach=False)
```

检查并统一所有可训练参数的设备与 dtype
初始化模型后，遍历检查所有参数的状态，确保统一：

for name, param in model.named_parameters():
    print(f"{name}: device={param.device}, dtype={param.dtype}")

若发现不匹配的参数，手动修正：

for param in model.parameters():
    param.data = param.data.to(device, dtype=torch.float32)

对齐PyTorch版本
本地与Colab的PyTorch版本差异可能导致兼容性问题，可在Colab中安装与本地一致的版本：
```
# 示例：安装指定版本
!pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2
```

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