PyTorch下,使用list放置模块,导致计算设备不一的报错

报错

在复现 Transformer 代码的训练阶段时,发生报错:

bash 复制代码
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

解决方案

通过next(linear.parameters()).device确定 model 已经在 cuda:0 上了,同时输入 model.forward()的张量也位于 cuda:0。输入的张量没什么好推敲的,于是考虑到模型具有多层结构,遂输出每层结构的设备信息,model.encoder -> model.encoder.sublayer[0] ··· ···

测试发现,model.encoder.sublayer[0] 之后的模块的设备信息均位于 cpu,原因是构造这部分模块时,由于需要多个相同的模块,使用了 list 来存放模块:

python 复制代码
# module: 需要深拷贝的模块
# n: 拷贝的次数
# return: 深拷贝后的模块列表
def clones(module, n: int) -> list:
    return [copy.deepcopy(module) for _ in range(n)]

显然 list 不支持 GPU,需要用 PyTorch 提供的代替:

python 复制代码
def clones(module, n: int):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])

ModuleList 把子模块存入列表,能像 Python 里普通的列表被索引,最重要的是能使内部的模块被正确注册,并对所有的 Module 方法可见。[Source]

成功解决!

相关环境

bash 复制代码
python                    3.11.7               he1021f5_0
pytorch                   2.1.2           py3.11_cuda12.1_cudnn8_0    
相关推荐
m0_7485548137 分钟前
golang如何实现用户订阅偏好管理_golang用户订阅偏好管理实现总结
jvm·数据库·python
RWKV元始智能41 分钟前
RWKV超并发项目教程,RWKV-LM训练提速40%
人工智能·rnn·深度学习·自然语言处理·开源
IT枫斗者1 小时前
前端部署后如何判断“页面是不是最新”?一套可落地的版本检测方案(适配 Vite/Vue/React/任意 SPA)
前端·javascript·vue.js·react.js·架构·bug
smj2302_796826521 小时前
解决leetcode第3911题.移除子数组元素后第k小偶数
数据结构·python·算法·leetcode
阿正呀2 小时前
Redis怎样实现本地缓存的高效失效通知
jvm·数据库·python
2501_901200532 小时前
mysql如何设置InnoDB引擎参数_优化innodb_buffer_pool
jvm·数据库·python
_.Switch2 小时前
东方财富股票数据JS逆向:secids字段和AES加密实战
开发语言·前端·javascript·网络·爬虫·python·ecmascript
AI技术增长2 小时前
Pytorch图像去噪实战(六):CBDNet真实噪声去噪实战,解决合成噪声模型落地效果差的问题
pytorch·深度学习·机器学习
Mr_sst2 小时前
Claude Code 部署与使用保姆级教程(2026 最新)
python·ai
瞎某某Blinder2 小时前
DFT学习记录[6]基于 HES06的能带计算+有效质量计算
python·学习·程序人生·数据挖掘·云计算·学习方法