Memory leak from evaluate() when using custom training loop and GPU

I am using multiple GPU's and AI-Platform for training my model, with a custom training loop in order to speed up the validation between epochs. During each epoch the memory slowly builds up, until OOM-exception is thrown a couple of epochs into training. Code for reference:

devices = tf.config.list_logical_devices('GPU')
strategy = tf.distribute.MirroredStrategy(devices)
with strategy.scope():
    model = create_model(...)
    model.compile(optimizer=tf.keras.optimizers.Adam(LEARNING_RATE))
    
    min_loss = 9999999999 # Used for early stopping
    for epoch in range(N_EPOCHS):
        start_training = time.time()
        model.fit(train.batch(BATCH_SIZE).cache())

        model.retrieval_task.factorized_metrics = (
            tfrs.metrics.FactorizedTopK(
                candidates=tfrs.layers.factorized_top_k.BruteForce().index_from_dataset(
                    items_ds.batch(1024).map(lambda item: (item["item_no"], model.item_model(item)))
                )
            )
        )
        model.compile(optimizer=tf.keras.optimizers.Adam(LEARNING_RATE))

        val_metrics = model.evaluate(val.batch(BATCH_SIZE), return_dict=True)
        loss = val_metrics['total_loss']

        # Early stopping, tracking total validation loss
        if loss < min_loss:
            best_model = model
            min_loss = loss
            patience_counter = 0
        else:
            patience_counter += 1
            if patience_counter > PATIENCE:
                model = best_model
                break

After around 4 epochs, during the .evaulate(...) part, the GPU's run out of memory:
"Allocator (GPU_0_bfc) ran out of memory trying to allocate 3.45GiB (rounded to 3708977408)requested by op retrieval_1/brute_force_2/TopKV2"

Any idea why, and what I can do to prevent this?