OutOfRangeError регистрируется в каждую эпоху после обновления с Tensorflow 1.8.0 до 1.11.0

У меня есть TfRecordsDataset:

...
dataset = dataset.batch(batch_size)
dataset = dataset.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
dataset_iterator = dataset.make_initializable_iterator()
(...) = dataset_iterator.get_next()

И тогда стандарт:

test_session.run(dataset_iterator.initializer)
while not done:
    try:
        ... = test_session.run([...])
        batch_count += 1
    except tf.errors.OutOfRangeError:
        done = True

Это отлично работало с Tensorflow 1.8.0. Сегодня я обновил Tensorflow до TF 11.1.0, и теперь у меня есть ошибка: OutOfRangeError: End of sequence [[{{node dataset/IteratorGetNext}} = IteratorGetNext

на каждом конце эпохи. Сама ошибка фактически поймана, но ошибка выводится на консоль. Может ли кто-нибудь указать, что на самом деле происходит?