Since `MAX_BATCH_PREFILL_TOKENS` is internally set by Text Generation
Inference (TGI) to `MAX_INPUT_TOKENS + 50`, and as the TGI on TPU model
warm-up validates that `MAX_BATCH_PREFILL_TOKENS <= MAX_INPUT_TOKENS *
BATCH_SIZE`, then we set the `BATCH_SIZE=2` so that `MAX_INPUT_TOKENS +
50 < MAX_INPUT_TOKENS * 2` so that the validation passes. Alternatively,
one could also set the `MAX_BATCH_PREFILL_TOKENS` to a value lower or
equal than `MAX_INPUT_TOKENS` (ideally equal).