多卡:关于trainer中线性层无法自动适配的小踩坑
transformer家的trainer适配了fp16、多卡...等,用起来很方便,但是之前我在实现模型的时候,发现在init里定义的线性层,在多卡上跑了以后,报错device不匹配,研究了许久,发现是trainer迁移的网络,不能被字典所包裹,最好直接定义为成员函数。例子如下:
在这里,self.linear.weight.device
的是对的,self.dict_['l'].weight.device
的device是不对的。