Tensorflow se queda sin memoria en el modelo.fit con preprocesamiento previo del conjunto de validación en la GPU.

Estoy jugando con este Colab localmente con una rtx 3070 de 8GB en Fedora 35 y tensorflow 2.4.0:
https://github.com/tensorflow/similarity/blob/master/examples/kaggle.ipynb
Tengo los mismos errores consistentes en Windows con una Nvidia GTX 1050 Ti (4GB).

Intenté desacoplar el origen del error OOM en model.fit y parece estar relacionado con la fase de preprocesamiento en la que estoy redimensionando el conjunto de validación. En esta fase se asigna memoria VRAM de la GPU.

<h1>Cargar imagen de validación en memoria</h1>

x_test = []
with tf.device('/CPU:0'): #solución para evitar ocupar memoria GPU y OOM en model.fit
    for p in tqdm(x_test_p):
        img = tf.io.read_file(p)
        img = tf.io.decode_image(img, dtype=tf.dtypes.float32)
        img = tf.image.resize_with_pad(img, IMG_SIZE, IMG_SIZE)
        #si es en escala de grises, convertir a RGB
        if tf.shape(img)[2]==3:
            pass
        else:
            img = tf.image.grayscale_to_rgb(img)
        x_test.append(img)

1. Si reduzco mucho el tamaño del conjunto de validación, model.fit tendrá éxito.
2. Si proceso toda la validación con la CPU en lugar de la GPU, model.fit tendrá éxito.
3. Si no paso todo el conjunto de validación preprocesado a model.fit, el error OOM seguirá presente. Es por eso que sugiero que es un problema relacionado únicamente con el preprocesamiento que ocupa memoria GPU útil.

¿Es posible que este preprocesamiento se cargue en VRAM y no se libere, limitando así el espacio de memoria GPU disponible para model.fit?

El problema es de alguna manera similar a esta pregunta de la que tomé la idea de preprocesar la validación con la CPU:
https://stackoverflow.com/questions/62571811/keras-oom-for-data-validation-using-gpu