最近在复现一个Transformer的小项目，用的是PyTorch 2.1。训练到大概第3个epoch的时候loss突然变成NaN，试过调小学习率、加梯度裁剪，还是会出现。同样的模型结构和数据，我用TensorFlow 2.15跑了一遍，居然稳定收敛了。我检查过输入数据没有脏值，也试过把混合精度关掉，依然会炸。想问下大家，这到底是框架差异还是我PyTorch的写法有坑？比如optimizer.zero_grad()的位置或者loss.backward()的调用时机？有点怀疑人生了，求指点。