Refactor assertions that use walrus

src/lightning/fabric/strategies/xla_fsdp.py

@@ -413,9 +413,10 @@ def save_checkpoint(
         # ensure model parameters are updated
         xm.mark_step()
 
+        parallel_devices = self.parallel_devices
+        assert parallel_devices is not None
         if self._sequential_save:
             # each host runs this in parallel, but the ranks in the host run it sequentially
-            assert (parallel_devices := self.parallel_devices) is not None
             for rank in range(len(parallel_devices)):
                 if rank == self.local_rank:
                     self._save_checkpoint_shard(path, state, storage_options, filter)
@@ -426,7 +427,6 @@ def save_checkpoint(
         if self._state_dict_type == "full":
             ckpt_prefix = str(path / "checkpoint")
             ckpt_suffix = "_rank-*-of-*.pth"
-            assert (parallel_devices := self.parallel_devices) is not None
             if len(parallel_devices) != self.world_size:  # multihost
                 raise OSError(
                     "Multihost setups do not have a shared filesystem, so the checkpoint shards cannot be consolidated"

src/lightning/pytorch/loops/evaluation_loop.py

@@ -122,7 +122,8 @@ def run(self) -> List[_OUT_DICT]:
             return []
         self.reset()
         self.on_run_start()
-        assert (data_fetcher := self._data_fetcher) is not None
+        data_fetcher = self._data_fetcher
+        assert data_fetcher is not None
         previous_dataloader_idx = 0
         while True:
             try:
@@ -374,7 +375,8 @@ def _evaluation_step(
 
         """
         trainer = self.trainer
-        assert (data_fetcher := self._data_fetcher) is not None
+        data_fetcher = self._data_fetcher
+        assert data_fetcher is not None
 
         if not (using_dataloader_iter := isinstance(data_fetcher, _DataLoaderIterDataFetcher)):
             batch = trainer.precision_plugin.convert_input(batch)

src/lightning/pytorch/loops/fetchers.py

@@ -52,10 +52,10 @@ def __iter__(self) -> "_DataFetcher":
         return self
 
     def __next__(self) -> _ITERATOR_RETURN:
-        assert (iterator := self.iterator) is not None
+        assert self.iterator is not None
         self._start_profiler()
         try:
-            batch = next(iterator)
+            batch = next(self.iterator)
         except StopIteration:
             self.done = True
             raise

src/lightning/pytorch/loops/fit_loop.py

@@ -337,13 +337,15 @@ def advance(self) -> None:
         """Runs one whole epoch."""
         log.debug(f"{type(self).__name__}: advancing loop")
 
-        assert (combined_loader := self._combined_loader) is not None
+        combined_loader = self._combined_loader
+        assert combined_loader is not None
         if combined_loader._mode == "sequential":
             raise ValueError(
                 f'`{type(self).__name__}` does not support the `CombinedLoader(mode="sequential")` mode.'
                 f" The available modes are: {[m for m in _SUPPORTED_MODES if m != 'sequential']}"
             )
-        assert (data_fetcher := self._data_fetcher) is not None
+        data_fetcher = self._data_fetcher
+        assert self._data_fetcher is not None
         data_fetcher.setup(combined_loader)
         with self.trainer.profiler.profile("run_training_epoch"):
             self.epoch_loop.run(data_fetcher)

src/lightning/pytorch/loops/prediction_loop.py

@@ -103,7 +103,8 @@ def run(self) -> Optional[_PREDICT_OUTPUT]:
             return None
         self.reset()
         self.on_run_start()
-        assert (data_fetcher := self._data_fetcher) is not None
+        data_fetcher = self._data_fetcher
+        assert data_fetcher is not None
         while True:
             try:
                 if isinstance(data_fetcher, _DataLoaderIterDataFetcher):
@@ -218,7 +219,8 @@ def _predict_step(
 
         """
         trainer = self.trainer
-        assert (data_fetcher := self._data_fetcher) is not None
+        data_fetcher = self._data_fetcher
+        assert data_fetcher is not None
 
         if not (using_dataloader_iter := isinstance(data_fetcher, _DataLoaderIterDataFetcher)):
             batch = trainer.precision_plugin.convert_input(batch)