From eb9c9b7f0d1c68cce3ae4e5fbf66a7b5254bf20f Mon Sep 17 00:00:00 2001 From: otaj Date: Wed, 28 Sep 2022 12:20:43 +0200 Subject: [PATCH] import from PL --- src/lightning_app/components/training.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/lightning_app/components/training.py b/src/lightning_app/components/training.py index aab4f8cdeb..b467615cca 100644 --- a/src/lightning_app/components/training.py +++ b/src/lightning_app/components/training.py @@ -42,7 +42,7 @@ class PyTorchLightningScriptRunner(TracerPythonScript): self.env = env def configure_tracer(self): - from lightning import Trainer + from pytorch_lightning import Trainer tracer = super().configure_tracer() tracer.add_traced(Trainer, "__init__", pre_fn=self._trainer_init_pre_middleware) @@ -70,8 +70,8 @@ class PyTorchLightningScriptRunner(TracerPythonScript): return super().run(**kwargs) def on_after_run(self, script_globals): - from lightning import Trainer - from lightning.pytorch.cli import LightningCLI + from pytorch_lightning import Trainer + from pytorch_lightning.cli import LightningCLI for v in script_globals.values(): if isinstance(v, LightningCLI):