Simplify rank_zero_experiment for torch.compile support (#17216)

Author: carmocca

src/lightning/fabric/loggers/logger.py

@@ -99,7 +99,7 @@ def rank_zero_experiment(fn: Callable) -> Callable:
     """Returns the real experiment on rank 0 and otherwise the _DummyExperiment."""
 
     @wraps(fn)
-    def experiment(self) -> Union[Any, _DummyExperiment]:  # type: ignore[no-untyped-def]
+    def experiment(self: Logger) -> Union[Any, _DummyExperiment]:
         """
         Note:
             ``self`` is a custom logger instance. The loggers typically wrap an ``experiment`` method
@@ -109,12 +109,9 @@ def experiment(self) -> Union[Any, _DummyExperiment]:  # type: ignore[no-untyped
             types that are specific to the custom logger. The return type here can be considered as
             ``Union[return type of logger.experiment, _DummyExperiment]``.
         """
-
-        @rank_zero_only
-        def get_experiment() -> Callable:
-            return fn(self)
-
-        return get_experiment() or _DummyExperiment()
+        if rank_zero_only.rank > 0:
+            return _DummyExperiment()
+        return fn(self)
 
     return experiment
 

src/lightning/pytorch/CHANGELOG.md

@@ -18,6 +18,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed WandbLogger not showing "best" aliases for model checkpoints when `ModelCheckpoint(save_top_k>0)` is used ([#17121](https://github.com/Lightning-AI/lightning/pull/17121))
 - Fixed the availability check for `rich` that prevented Lightning to be imported in Google Colab ([#17156](https://github.com/Lightning-AI/lightning/pull/17156))
 - Fixed parsing the precision config for inference in `DeepSpeedStrategy` ([#16973](https://github.com/Lightning-AI/lightning/pull/16973))
+- Fixed issue where `torch.compile` would fail when logging to WandB ([#17216](https://github.com/Lightning-AI/lightning/pull/17216))
 
 
 ## [2.0.0] - 2023-03-15