训练GNN时首个训练样本后Loss出现NaN问题求助

核心问题 #

训练GNN时，首个epoch返回有限Loss，但第二个epochLoss立即变为NaN。未训练模型对前5个输入的预测Loss为有限值，且因输入处于带周期性边界条件的3D空间，Loss受最小镜像约束在0到半盒长之间，希望避免数据缩放/归一化。

输入数据说明 #

模型接收5类输入：

• 长度为385的节点标签向量（取值为1、2或3）
• 385×5×300的tensor：385个节点各对应5条边，每条边包含300维特征向量
• 385×5的节点索引矩阵：记录每个节点的连接索引
• 384×3的矩阵：代表1类和2类节点的3D坐标
• 1×3的向量：代表单个3类节点的3D坐标

任务目标 #

仅需对指定的3类节点进行预测，输出其3D位置的更新值。

已提供的相关内容 #

已附上以下用于复现和排查的内容：

• 输入i1-i5与目标ys_list的定义
• 合成数据的复现方法
• 模型类代码
• 相关辅助函数
• 训练流程代码
• 损失函数及回调函数

内容的提问来源于stack exchange，提问作者Will Borrelli