# Batch size for trainer.predict()

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374%3E%5C
**Category:** 🤗Transformers\
**Created:** [January 25, 2021, 3:21pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374 "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](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fberkayberabi)%5C
**Post date:** [January 25, 2021, 3:21pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374%2F1 "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](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsgugger)%5C
**Post date:** [January 25, 2021, 4:27pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374%2F2 "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](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmhilmiasyrofi)%5C
**Post date:** [August 4, 2021, 5:25am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374%2F3 "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](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fkechan)%5C
**Post date:** [July 17, 2022, 10:19pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374%2F4 "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](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsriram6399)%5C
**Post date:** [November 26, 2022, 4:08pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbatch-size-for-trainer-predict%2F3374%2F5 "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.