Keras实现论文对抗训练方法时计算损失出现NaN值问题求助

我现在正在复现一篇论文里的对抗训练方法，核心思路是沿着网络对微小变化最敏感的方向修改输入数据，然后用修改后的数据带着原标签来训练网络。

论文里的关键定义和参数如下：

• d：单位向量
• xhi：非零数
• D：损失函数，本文采用稀疏交叉熵
• 训练损失：结合带标签数据的损失度量l与算法1中的梯度项Rvadv，论文选择alpha=1，以此将模型在带标签数据集上的性能纳入损失计算

我用Keras结合MNIST数据集、100大小的小批量来实现这个方法，但在执行梯度下降更新权重的过程中，只迭代了几次就出现了NaN值，现在完全定位不到问题根源。我已经把代码放到了Colab会话和Gist里。

有没有大佬能帮忙排查下NaN出现的原因？

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