From bdd690a74da5283cbc893dfd79e1c7c72ec1bcfa Mon Sep 17 00:00:00 2001
From: yujun <50394665+JunnYu@users.noreply.github.com>
Date: Mon, 2 May 2022 19:49:19 +0800
Subject: [PATCH] add torch.no_grad when in eval mode (#17020)

* add torch.no_grad when in eval mode

* make style quality
---
 .../run_image_classification_no_trainer.py                    | 3 ++-
 .../run_semantic_segmentation_no_trainer.py                   | 3 ++-
 examples/pytorch/text-classification/run_glue_no_trainer.py   | 4 +++-
 .../run_{{cookiecutter.example_shortcut}}.py                  | 4 +++-
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/examples/pytorch/image-classification/run_image_classification_no_trainer.py b/examples/pytorch/image-classification/run_image_classification_no_trainer.py
index 39f805b458..64f4601f3a 100644
--- a/examples/pytorch/image-classification/run_image_classification_no_trainer.py
+++ b/examples/pytorch/image-classification/run_image_classification_no_trainer.py
@@ -469,7 +469,8 @@ def main():
         model.eval()
         samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
-            outputs = model(**batch)
+            with torch.no_grad():
+                outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             predictions, references = accelerator.gather((predictions, batch["labels"]))
             # If we are in a multiprocess environment, the last batch has duplicates
diff --git a/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py b/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
index d5a6a16fe4..77fc95b2a6 100644
--- a/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
+++ b/examples/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py
@@ -579,7 +579,8 @@ def main():
         model.eval()
         samples_seen = 0
         for step, batch in enumerate(tqdm(eval_dataloader, disable=not accelerator.is_local_main_process)):
-            outputs = model(**batch)
+            with torch.no_grad():
+                outputs = model(**batch)
 
             upsampled_logits = torch.nn.functional.interpolate(
                 outputs.logits, size=batch["labels"].shape[-2:], mode="bilinear", align_corners=False
diff --git a/examples/pytorch/text-classification/run_glue_no_trainer.py b/examples/pytorch/text-classification/run_glue_no_trainer.py
index 73f52825a3..860e4035a3 100644
--- a/examples/pytorch/text-classification/run_glue_no_trainer.py
+++ b/examples/pytorch/text-classification/run_glue_no_trainer.py
@@ -22,6 +22,7 @@ import random
 from pathlib import Path
 
 import datasets
+import torch
 from datasets import load_dataset, load_metric
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
@@ -514,7 +515,8 @@ def main():
         model.eval()
         samples_seen = 0
         for step, batch in enumerate(eval_dataloader):
-            outputs = model(**batch)
+            with torch.no_grad():
+                outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1) if not is_regression else outputs.logits.squeeze()
             predictions, references = accelerator.gather((predictions, batch["labels"]))
             # If we are in a multiprocess environment, the last batch has duplicates
diff --git a/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py b/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
index 0d9a0c6d32..5a641f85f2 100755
--- a/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
+++ b/templates/adding_a_new_example_script/{{cookiecutter.directory_name}}/run_{{cookiecutter.example_shortcut}}.py
@@ -28,6 +28,7 @@ from dataclasses import dataclass, field
 from typing import Optional, List
 
 import datasets
+import torch
 from datasets import load_dataset
 
 import transformers
@@ -871,7 +872,8 @@ def main():
 
         model.eval()
         for step, batch in enumerate(eval_dataloader):
-            outputs = model(**batch)
+            with torch.no_grad(): 
+                outputs = model(**batch)
             predictions = outputs.logits.argmax(dim=-1)
             metric.add_batch(
                 predictions=accelerator.gather(predictions),