Tf.data API не может распечатать все партии

Я самоучусь по tf.data API. Я использую набор данных MNIST для двоичной классификации. Данные для обучения x и y заархивированы вместе в полном наборе train_dataset. Вместе с этим методом zip сначала связан метод набора данных batch(). данные сгруппированы с размером пакета 30. Поскольку размер моего обучающего набора составляет 11623, с размером пакета 128 у меня будет 91 пакет. Размер последней партии будет 103, и это нормально, так как это LSTM. Дополнительно использую выпадение. Когда я вычисляю точность партии, я отключаю выпадение.

Полный код приведен ниже:

#Ignore the warnings
import warnings
warnings.filterwarnings("ignore")

import pandas as pd
import tensorflow as tf
import numpy as np

import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (8,7)

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/")

Xtrain = mnist.train.images[mnist.train.labels < 2]
ytrain = mnist.train.labels[mnist.train.labels < 2]

print(Xtrain.shape)
print(ytrain.shape)

#Data parameters
num_inputs = 28
num_classes = 2
num_steps=28

# create the training dataset
Xtrain = tf.data.Dataset.from_tensor_slices(Xtrain).map(lambda x: tf.reshape(x,(num_steps, num_inputs)))
# apply a one-hot transformation to each label for use in the neural network
ytrain = tf.data.Dataset.from_tensor_slices(ytrain).map(lambda z: tf.one_hot(z, num_classes))
# zip the x and y training data together and batch and Prefetch data for faster consumption
train_dataset = tf.data.Dataset.zip((Xtrain, ytrain)).batch(128).prefetch(128)

iterator = tf.data.Iterator.from_structure(train_dataset.output_types,train_dataset.output_shapes)
X, y = iterator.get_next()

training_init_op = iterator.make_initializer(train_dataset)


#### model is here ####

#Network parameters
num_epochs = 2
batch_size = 128
output_keep_var = 0.5

with tf.Session() as sess:
    init.run()

    print("Initialized")
    # Training cycle
    for epoch in range(0, num_epochs):
        num_batch = 0
        print ("Epoch: ", epoch)
        avg_cost = 0.
        avg_accuracy =0
        total_batch = int(11623 / batch_size + 1)
        sess.run(training_init_op)
       while True:
            try:
                _, miniBatchCost = sess.run([trainer, loss], feed_dict = {output_keep_prob: output_keep_var})
                miniBatchAccuracy = sess.run(accuracy, feed_dict = {output_keep_prob: 1.0})
               print('Batch %d: loss = %.2f, acc = %.2f' % (num_batch, miniBatchCost, miniBatchAccuracy * 100))
                num_batch +=1
            except tf.errors.OutOfRangeError:
                break

Когда я запускаю этот код, кажется, что он работает и печатает:

Batch 0: loss = 0.67276, acc = 0.94531
Batch 1: loss = 0.65672, acc = 0.92969
Batch 2: loss = 0.65927, acc = 0.89062
Batch 3: loss = 0.63996, acc = 0.99219
Batch 4: loss = 0.63693, acc = 0.99219
Batch 5: loss = 0.62714, acc = 0.9765
......
......
Batch 39: loss = 0.16812, acc = 0.98438
Batch 40: loss = 0.10677, acc = 0.96875
Batch 41: loss = 0.11704, acc = 0.99219
Batch 42: loss = 0.10592, acc = 0.98438
Batch 43: loss = 0.09682, acc = 0.97656
Batch 44: loss = 0.16449, acc = 1.00000

Однако, как легко увидеть, что-то не так. Напечатано всего 45 партий, а не 91, и я не знаю, почему это происходит. Я так много пробовал и думаю, что что-то упускаю.

Я могу использовать функцию repeat(), но я не хочу этого, потому что у меня есть избыточные наблюдения для последних пакетов, и я хочу, чтобы LSTM справился с этим.