HPU & TPU doesn't support torch.inference_mode

CHANGELOG.md

@@ -185,7 +185,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed mismatching default values for the types of some arguments in the DeepSpeed and Fully-Sharded strategies which made the CLI unable to use them ([#12989](https://github.com/PyTorchLightning/pytorch-lightning/pull/12989))
 
 
--
+- Fixed an issue with unsupported torch.inference_mode() on hpu backends by making it use no_grad ([#13014](https://github.com/PyTorchLightning/pytorch-lightning/pull/13014))
 
 
 -

docs/source/accelerators/hpu_basic.rst

@@ -81,3 +81,4 @@ Known limitations
 * Multiple optimizers are not supported.
 * `Habana dataloader <https://docs.habana.ai/en/latest/PyTorch_User_Guide/PyTorch_User_Guide.html#habana-data-loader>`__ is not supported.
 * :class:`~pytorch_lightning.callbacks.device_stats_monitor.DeviceStatsMonitor` is not supported.
+* torch.inference_mode() is not supported

pytorch_lightning/trainer/trainer.py

@@ -2802,10 +2802,12 @@ def configure_optimizers(self):
 
 @contextmanager
 def _evaluation_context() -> Generator:
-    # inference mode is not supported with gloo backend (#9431)
+    # inference mode is not supported with gloo backend (#9431) and hpu backend
     context_manager_class = (
         torch.inference_mode
-        if _TORCH_GREATER_EQUAL_1_9 and not (dist.is_initialized() and dist.get_backend() == "gloo")
+        if _TORCH_GREATER_EQUAL_1_9
+        and not (dist.is_initialized() and dist.get_backend() == "gloo")
+        and not _HPU_AVAILABLE
         else torch.no_grad
     )
     with context_manager_class():