[BUG] CUDA Out of Memory when eval model.

[Crystalxd]

Required prerequisites

• I have read the documentation https://github.com/baichuan-inc/baichuan-7B/blob/HEAD/README.md.
• I have searched the Issue Tracker and Discussions that this hasn't already been reported. (+1 or comment there if it has.)
• Consider asking first in a Discussion.

System information

conda environment
torch=2.0.1
transformers=4.29.2
...

Problem description

I used A100(80G) to run the evaluate_zh.py script for evaluating baichuan model, but it occupied abundant GPU memory up to overflow. Then I found the model loaded without eval mode, meanwhile, it inferred without no_grad.

Reproducible example code

The Python snippets:

[https://github.com/baichuan-inc/Baichuan-7B/blob/6f3ef4633a90c2d8a3e0763d0dec1b8dc11588f5/evaluation/evaluate_zh.py#L97C13-L97C13](url)
self.model = model.eval()

https://github.com/baichuan-inc/Baichuan-7B/blob/6f3ef4633a90c2d8a3e0763d0dec1b8dc11588f5/evaluation/evaluate_zh.py#L103
Add on this line:
@torch.inference_mode()

Command lines:

Extra dependencies:

Steps to reproduce:

Traceback

No response

Expected behavior

No response

Additional context

No response

Checklist

• I have provided all relevant and necessary information above.
• I have chosen a suitable title for this issue.

[Guanze-Chen]

Thank you. It works!!!

[ICanFlyGFC]

Thanks!

[Guanze-Chen]

您的邮件已经收到，会尽快回复您