多GPU训练with attention_mask transformers/modeling_gpt2.py line 403 StopIteration

[Yuchen-Li]

版本：

python 3.6.0
torch 1.6.0+cu101

在多GPU without attention_mask的时候一切正常， 但是加上attention_mask后在forward阶段报错如下:

File "$PATH/envs/gpt_venv/lib/python3.6/site-packages/transformers/modeling_gpt2.py", line 403, in forward
    attention_mask = attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
StopIteration

请问有没有人遇到过相似的错误呢？

[Yuchen-Li]

debug了一下尝试了不同的torch版本，torch==1.4, 1.5均会报错Segmentation fault (core dumped)