From 9bfd40beb121efdd684d3bec397c69e63c74c3c8 Mon Sep 17 00:00:00 2001 From: awaelchli Date: Thu, 29 Sep 2022 14:49:02 +0200 Subject: [PATCH] move it as early as possible --- src/pytorch_lightning/trainer/trainer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/pytorch_lightning/trainer/trainer.py b/src/pytorch_lightning/trainer/trainer.py index c899ec3d059a0..104241eacc8fc 100644 --- a/src/pytorch_lightning/trainer/trainer.py +++ b/src/pytorch_lightning/trainer/trainer.py @@ -951,6 +951,9 @@ def _restore_modules_and_callbacks(self, checkpoint_path: Optional[_PATH] = None def _run( self, model: "pl.LightningModule", ckpt_path: Optional[str] = None ) -> Optional[Union[_EVALUATE_OUTPUT, _PREDICT_OUTPUT]]: + + self.strategy.setup_environment() + if self.state.fn in (TrainerFn.FITTING, TrainerFn.TUNING): min_epochs, max_epochs = _parse_loop_limits( self.min_steps, self.max_steps, self.min_epochs, self.max_epochs, self @@ -979,7 +982,7 @@ def _run( # ---------------------------- self._call_callback_hooks("on_before_accelerator_backend_setup") log.detail(f"{self.__class__.__name__}: setting up strategy environment") - self.strategy.setup_environment() + self.__setup_profiler() self._call_setup_hook() # allow user to setup lightning_module in accelerator environment