# Batch size for trainer.predict()

**URL:** <https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374>\
**Category:** 🤗Transformers\
**Created:** [January 25, 2021, 3:21pm UTC](https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374 "2021-01-25T15:21:38Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![berkayberabi](https://avatars.discourse-cdn.com/v4/letter/b/258eb7/32.png) [@berkayberabi](https://discuss.huggingface.co/u/berkayberabi)\
**Post date:** [January 25, 2021, 3:21pm UTC](https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374/1 "2021-01-25T15:21:38Z")

</div>

Hi,

I pass a test dataset to trainer.predict but I have many samples. Therefore, I get a memory error. Does the library support a way of batch based trainer.predict? or do I have to implement it myself?

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [January 25, 2021, 4:27pm UTC](https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374/2 "2021-01-25T16:27:33Z")

</div>

You can pass `eval_accumulation_steps=xxx` to pass the predictions to the CPU every `xxx` steps, this should help.

---

<div class="post-metadata">

**Author:** ![mhilmiasyrofi](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mhilmiasyrofi/32/2442_2.png) [@mhilmiasyrofi](https://discuss.huggingface.co/u/mhilmiasyrofi)\
**Post date:** [August 4, 2021, 5:25am UTC](https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374/3 "2021-08-04T05:25:18Z")

</div>

You can set the batch size manually using trainer.prediction\_loop()

Instead of using `trainer.predict(test_dataset)`, you can use torch DataLoader for trainer.prediction\_loop(). Thus, you might change

from

```auto
raw_pred, _, _ = trainer.predict(test_dataset)

```

into:

```auto
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
raw_pred, _, _ = trainer.prediction_loop(test_loader, description="prediction")

```

---

<div class="post-metadata">

**Author:** ![kechan](https://avatars.discourse-cdn.com/v4/letter/k/aca169/32.png) [@kechan](https://discuss.huggingface.co/u/kechan)\
**Post date:** [July 17, 2022, 10:19pm UTC](https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374/4 "2022-07-17T22:19:18Z")

</div>

In transformers 4.20.1,

> args = TrainingArguments(output\_dir=‘tmp\_trainer’, per\_device\_eval\_batch\_size=16)
> 
> trainer = Trainer(model=model, args=args)
> 
> predictions = trainer.predict(pred\_dataset)

---

<div class="post-metadata">

**Author:** ![sriram6399](https://avatars.discourse-cdn.com/v4/letter/s/258eb7/32.png) [@sriram6399](https://discuss.huggingface.co/u/sriram6399)\
**Post date:** [November 26, 2022, 4:08pm UTC](https://discuss.huggingface.co/t/batch-size-for-trainer-predict/3374/5 "2022-11-26T16:08:11Z")

</div>

Hi I tried this method, but I see that the prediction process is killed at 99% without generating the predictions. There are no Memory Issues. Looks like if I use trainer.prediction\_loop() method, I cannot set the argument predict\_with\_gentrate=True, I am thinking this might be causing the problem, I am not sure though. I am very new to working on pre\_trained models. Could you please let me know your insights, what could be the possible reason for this issue.