Fix

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

@@ -49,6 +49,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - 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 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