Fix inference mode grad context manager

src/pytorch_lightning/trainer/trainer.py

@@ -2232,7 +2232,7 @@ def _evaluation_context(accelerator: Accelerator, inference_mode: bool = True) -
         and not (dist.is_available() and dist.is_initialized() and dist.get_backend() == "gloo")
         and not isinstance(accelerator, HPUAccelerator)
         and not isinstance(accelerator, TPUAccelerator)
-        else torch.no_grad
+        else torch.enable_grad
     )
     with context_manager_class():
         yield

tests/tests_pytorch/trainer/flags/test_inference_mode.py

@@ -21,10 +21,16 @@
 def test_eval_inference_mode():
     """Testing overwriting trainer arguments."""
 
+    class BoringModelGrad(BoringModel):
+        def on_test_epoch_start(self) -> None:
+            assert torch.is_grad_enabled()
+            assert not torch.is_inference_mode_enabled()
+            return super().on_test_epoch_start()
+
     class BoringModelNoGrad(BoringModel):
         def on_test_epoch_start(self) -> None:
             assert not torch.is_grad_enabled()
-            assert not torch.is_inference_mode_enabled()
+            assert torch.is_inference_mode_enabled()
             return super().on_test_epoch_start()
 
     class BoringModelForInferenceMode(BoringModel):
@@ -34,6 +40,8 @@ def on_test_epoch_start(self) -> None:
             return super().on_test_epoch_start()
 
     trainer = Trainer(logger=False, inference_mode=False, fast_dev_run=True)
+    trainer.test(BoringModelGrad())
+    trainer = Trainer(logger=False, inference_mode=True, fast_dev_run=True)
     trainer.test(BoringModelNoGrad())
     trainer = Trainer(logger=False, inference_mode=True, fast_dev_run=True)
     trainer.test(BoringModelForInferenceMode())