move it as early as possible

Lightning-AI · Sep 29, 2022 · 9bfd40b · 9bfd40b
1 parent eb69b9e
commit 9bfd40b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git 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