From bb14a979e579baaa301ceccba42190c8d7bbf989 Mon Sep 17 00:00:00 2001 From: Jerry Mannil Date: Wed, 10 Jan 2024 11:05:48 -0800 Subject: [PATCH] Call `prepare_data()` after `setup_environment()` for XLA (#19181) Co-authored-by: awaelchli --- src/lightning/pytorch/trainer/trainer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/lightning/pytorch/trainer/trainer.py b/src/lightning/pytorch/trainer/trainer.py index cc335c975e918..592bb7196519a 100644 --- a/src/lightning/pytorch/trainer/trainer.py +++ b/src/lightning/pytorch/trainer/trainer.py @@ -935,10 +935,6 @@ def _run( _verify_loop_configurations(self) - # hook - log.debug(f"{self.__class__.__name__}: preparing data") - self._data_connector.prepare_data() - # ---------------------------- # SET UP THE TRAINER # ---------------------------- @@ -946,6 +942,9 @@ def _run( self.strategy.setup_environment() self.__setup_profiler() + log.debug(f"{self.__class__.__name__}: preparing data") + self._data_connector.prepare_data() + call._call_setup_hook(self) # allow user to set up LightningModule in accelerator environment log.debug(f"{self.__class__.__name__}: configuring model") call._call_configure_model(self)