我在使用GPU和TPU的谷歌colab训练中遇到了一个相当奇怪的问题，我使用了自定义损失，GPU上的tfrecord数据集也可以，但如果我切换到TPU，则会将nan作为验证损失。没有其他特定错误。此外，旧的验证tfrecord数据集在TPU上运行良好。这让我觉得数据可能有一些特定的东西。这也出现在model.eevaluate(…(上，因为它在验证集上。
知道如何最好地用TPU调试这个吗？可根据要求提供更多详细信息。
我的问题可能与https://github.com/tensorflow/tensorflow/issues/41635(尽管它甚至可以用于非自定义损失函数(。就我的情况而言，我没有看到它使用开箱即用的丢失功能，但当我使用自定义丢失时，我看到了它。自定义丢失似乎不是主要原因，因为它在CPU和GPU下都适用于任何数据集。
无论如何，我遵循了这个问题的提示，丢弃了最后一批(它的大小小于batch_size(，NaN不再可见。虽然这解决了问题，但我仍然没有一个明确的答案来解决根本原因。
这是解决问题的解决方案：
在model.fit((中添加这些参数，就会正常工作
steps_per_epoch = train.shape[0]//batch_Size
validation_steps = validate.shape[0]//batch_Size
这将删除最后一批要处理的产品，然后问题将得到解决，不再有Nan出现