Merge pull request #2534 from GuillaumeErhard/inference_mode

Update no_grad usage to inference_mode if possible

Author: rwightman

benchmark.py

@@ -321,7 +321,7 @@ def _step():
             f'Running inference benchmark on {self.model_name} for {self.num_bench_iter} steps w/ '
             f'input size {self.input_size} and batch size {self.batch_size}.')
 
-        with torch.no_grad():
+        with torch.inference_mode():
             self._init_input()
 
             for _ in range(self.num_warm_iter):