Predicción del modelo de TensorFlow que resulta en una fuga de memoria (fuera de memoria)

Estoy tratando de ejecutar predicciones para un modelo de clasificación de imágenes en un gran número de archivos de imagen (800000). A continuación se muestra un fragmento de código de cómo se ve mi canalización de datos:

import tensorflow as tf
BATCH_SIZE = 128

def configure_for_performance(ds):
    ds = ds.batch(BATCH_SIZE)
    ds = ds.prefetch(buffer_size=AUTOTUNE)
    tf.keras.backend.clear_session()
    return ds

def decode_img(img):
    img = tf.io.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, [IMG_HEIGHT, IMG_WIDTH])
    return img

def process_path(file_path):
    img = tf.io.read_file(file_path)
    img = decode_img(img)
    return img

list_ds = tf.data.Dataset.list_files([filepath1,filepath2,...,filepathN],
                                     shuffle=False)

files_list = list()
for files in list_ds.as_numpy_iterator():
    files_list.append(files.decode("utf-8"))

test_ds = list_ds.map(process_path, num_parallel_calls=AUTOTUNE)
test_size = test_ds.cardinality().numpy()
test_ds = configure_for_performance(test_ds)
test_ds = test_ds.map(
    lambda images: normalization_sequential_layer(images))

test_steps = test_size // BATCH_SIZE + 1
score = tf.nn.softmax(model.predict(test_ds, steps=test_steps))

Lo estoy ejecutando en una máquina con 32 núcleos y 64 GB de RAM (sin GPU), pero el proceso tarda bastante tiempo y luego se cierra con el mensaje + Killed, lo que, según mi investigación, podría relacionarse con un problema de falta de RAM (fuga de memoria que ocurre con el tiempo y después de un punto se sale de la RAM). Según mi investigación adicional, puedo ver que las personas enfrentan problemas de OOM al tener un bucle for alrededor del modelo.predict, no estoy seguro de cuál podría ser la razón aquí.