可能的原因是，cpu用的是float32计算，而gpu可能使用amp后用float16计算导致结果溢出
看网络是否会过拟合，也就是loss变得非常低，而不是nan。如果有nan，或者loss不下降，代码必然有问题。
Batch越大越容易爆炸，建议调小batch