Hi @sgugger, I understood the purpose of predict_with_generate from the example script. Normally, the forward pass of the model returns loss and logits, but we need tokens for the ROUGE/BLEU, where generate() comes into picture and predict_with_generate handles that. But, then, I am not able to understand the cause for my RAM usage exceeding 100%? I tried caching the datasets, defined them globally, used xmp.SerialExecutor(). But, with every epoch RAM usages keeps increasing? I am not able to finetune t5-small for more than 4 epochs if I pass compute_metrics.