_fit_impl

Author: Borda

src/lightning/pytorch/trainer/trainer.py

@@ -562,8 +562,6 @@ def fit(
         _verify_strategy_supports_compile(model, self.strategy)
         self.state.fn = TrainerFn.FITTING
         self.state.status = TrainerStatus.RUNNING
-        if _is_registry(ckpt_path) and module_available("litmodels"):
-            download_model_from_registry(ckpt_path, self)
         self.training = True
         self.should_stop = False
         call._call_and_handle_interrupt(
@@ -580,6 +578,9 @@ def _fit_impl(
     ) -> None:
         log.debug(f"{self.__class__.__name__}: trainer fit stage")
 
+        if _is_registry(ckpt_path) and module_available("litmodels"):
+            download_model_from_registry(ckpt_path, self)
+
         # if a datamodule comes in as the second arg, then fix it for the user
         if isinstance(train_dataloaders, LightningDataModule):
             datamodule = train_dataloaders