Иметь одинаковый результат в двух разных запусках tensorflow

Я пытаюсь запустить пример в Обучение нейронной сети на MNIST с помощью Keras. Я хочу сделать это дважды, указав веса и не перемешивая, чтобы получить ОДИНАКОВЫЙ результат в обоих прогонах. Вот полный код:

import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
import numpy as np

tf.enable_v2_behavior()
tfds.disable_progress_bar()

def normalize_img(image, label):
    """Normalizes images: `uint8` -> `float32`."""
    return tf.cast(image, tf.float32) / 255., label

def get_dataset():
    (ds_train, ds_test), ds_info = tfds.load(
        'mnist',
        split=['train', 'test'],
        shuffle_files=True,
        as_supervised=True,
        with_info=True,
    )

    ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    ds_train = ds_train.cache()
    # ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
    ds_train = ds_train.batch(128)
    ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)

    ds_test = ds_test.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    ds_test = ds_test.batch(128)
    ds_test = ds_test.cache()
    ds_test = ds_test.prefetch(tf.data.experimental.AUTOTUNE)

    return ds_train, ds_test

def keras_fit(ds_train, ds_test, verbose=True, init1='glorot_uniform', init2='glorot_uniform'):
    # https://www.tensorflow.org/datasets/keras_example
    model = tf.keras.models.Sequential([
      tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
      tf.keras.layers.Dense(128, activation='relu', kernel_initializer=init1),
      tf.keras.layers.Dense(10, activation='softmax', kernel_initializer=init2)
    ])
    model.compile(
        loss='sparse_categorical_crossentropy',
        optimizer=tf.keras.optimizers.Adam(0.001),
        metrics=['accuracy'],
    )
    model.fit(
        ds_train,
        epochs=6,
        validation_data=ds_test,
        verbose=verbose, shuffle=False
    )
    return model.evaluate(ds_test, verbose=verbose)

def test_mnist():
    init = tf.keras.initializers.GlorotUniform()
    init1 = tf.constant_initializer(init((784, 128)).numpy())
    init2 = tf.constant_initializer(init((128, 10)).numpy())
    ds_train, ds_test = get_dataset()
    keras1 = keras_fit(ds_train, ds_test, init1=init1, init2=init2)
    keras2 = keras_fit(ds_train, ds_test, init1=init1, init2=init2)
    print(keras1)
    print(keras2)

if __name__ == "__main__":
    test_mnist()

Насколько я понимаю, оба слоя будут инициализированы с одинаковыми значениями, и я не перемешиваю данные случайным образом (функция подгонки имеет shuffle=False). Разве я не должен получить такой же точный результат? Что я делаю не так?

Я получаю довольно похожие результаты и иногда одинаковую точность, но это не на 100%.

PS: я получаю следующее сообщение после каждой эпохи:

[[{{node IteratorGetNext}}]]
         [[Shape/_6]]
2020-12-20 13:52:10.002034: W tensorflow/core/common_runtime/base_collective_executor.cc:217] BaseCollectiveExecutor::StartAbort Out of range: End of sequence
         [[{{node IteratorGetNext}}]]

Пробовал в принципе тот же скрипт, но с модом MNIST. И это сработало! Я получаю точно такой же результат.

Разница в том, что Fashion MNIST использует пустой массив, а не tf.dataset.

Проблема должна быть либо в случайном порядке, сделанном классом tf.dataset, либо в предупреждающем сообщении, которое я получаю, может быть, я не заканчиваю перебирать полный набор данных, прежде чем остановлюсь и начну следующий в той же точке?