TensorFlow训练时数据集提前耗尽（OUT_OF_RANGE）问题求助

大家好，我现在遇到了TensorFlow训练的问题，想请教下各位大佬。

我正在从一个包含图像路径和掩码路径（对应变量xtr和ytr）的CSV文件构建数据集，初始CSV大概有1000条记录。每张图像从云存储加载后，会根据自身高度用flat_map拆分成4张或8张小图，按说数据集的元素数至少不会低于初始的1000条。

我的数据处理管道代码如下：

train_ds = (
    tf.data.Dataset.from_tensor_slices((xtr, ytr))
    .shuffle(200)
    .map(self._load_image, num_parallel_calls=tf.data.AUTOTUNE)
    .map(self._preprocess, num_parallel_calls=tf.data.AUTOTUNE)
    .flat_map(self._split)
    .map(self._augment)
    .batch(self._batch_size, drop_remainder=True)
    .prefetch(tf.data.experimental.AUTOTUNE)
)

统计下来，拆分后的数据集大概有1700个元素。训练用的batch size是4，每个epoch我大概设置了250步。

但现在训练时频繁出现数据耗尽的错误，比如会弹出：

Local rendezvous is aborting with status: OUT_OF_RANGE

有时候还会收到明确的用户警告，提示数据已经耗尽了。

有没有大佬能给我一些解决这个问题的思路呀？非常感谢！

内容来源于stack exchange