Fixes for LightningCLI

Author: carmocca

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -228,7 +228,7 @@ def update_train_step_metrics(self) -> None:
             return
 
         # TODO: remove this call in v1.7
-        self._log_gpus_metrics()
+        # self._log_gpus_metrics()
 
         # when metrics should be logged
         assert not self._epoch_end_reached

pytorch_lightning/trainer/trainer.py

@@ -1415,7 +1415,7 @@ def call_hook(
                 # Rely on the accelerator output if lightningModule hook returns nothing
                 # Required for cases such as DataParallel where we reduce the output for the user
                 # todo: move this data parallel logic into the data parallel plugin
-                output = accelerator_output if output is None else output
+                output = ttp_output if output is None else output
 
             # call the ttp hook
             if hook_name not in ("setup", "teardown", "on_train_start") and hasattr(

tests/helpers/boring_model.py

@@ -158,18 +158,15 @@ def prepare_data(self):
     def setup(self, stage: Optional[str] = None):
         if stage == "fit" or stage is None:
             self.random_train = Subset(self.random_full, indices=range(64))
-            self.dims = self.random_train[0].shape
 
         if stage in ("fit", "validate") or stage is None:
             self.random_val = Subset(self.random_full, indices=range(64, 64 * 2))
 
         if stage == "test" or stage is None:
             self.random_test = Subset(self.random_full, indices=range(64 * 2, 64 * 3))
-            self.dims = getattr(self, "dims", self.random_test[0].shape)
 
         if stage == "predict" or stage is None:
             self.random_predict = Subset(self.random_full, indices=range(64 * 3, 64 * 4))
-            self.dims = getattr(self, "dims", self.random_predict[0].shape)
 
     def train_dataloader(self):
         return DataLoader(self.random_train)

tests/utilities/test_cli.py

@@ -574,7 +574,7 @@ def add_arguments_to_parser(self, parser):
 
 class EarlyExitTestModel(BoringModel):
     def on_fit_start(self):
-        raise Exception("Error on fit start")
+        raise MisconfigurationException("Error on fit start")
 
 
 @pytest.mark.parametrize("logger", (False, True))
@@ -586,8 +586,10 @@ def on_fit_start(self):
         pytest.param({"tpu_cores": 1}, marks=RunIf(tpu=True)),
     ),
 )
-def test_cli_ddp_spawn_save_config_callback(tmpdir, logger, trainer_kwargs):
-    with mock.patch("sys.argv", ["any.py", "fit"]), pytest.raises(Exception, match=r"Error on fit start"):
+def test_cli_distributed_save_config_callback(tmpdir, logger, trainer_kwargs):
+    with mock.patch("sys.argv", ["any.py", "fit"]), pytest.raises(
+        MisconfigurationException, match=r"Error on fit start"
+    ):
         LightningCLI(
             EarlyExitTestModel,
             trainer_defaults={