Fig logging with log_gpu_memory='min_max' (#9013)

Author: SkafteNicki

CHANGELOG.md

@@ -37,6 +37,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed a bug in the binary search mode of auto batch size scaling where exception was thrown if the first trainer run resulted in OOM ([#8954](https://github.com/PyTorchLightning/pytorch-lightning/pull/8954))
 
 
+- Fixed a bug causing logging with `log_gpu_memory='min_max'` not working ([#9013](https://github.com/PyTorchLightning/pytorch-lightning/pull/9013))
+
 ## [1.4.0] - 2021-07-27
 
 ### Added

pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -226,9 +226,14 @@ def update_train_epoch_metrics(self) -> None:
 
     def _log_gpus_metrics(self):
         for key, mem in self.gpus_metrics.items():
-            gpu_id = int(key.split("/")[0].split(":")[1])
-            if gpu_id in self.trainer.accelerator_connector.parallel_device_ids:
-                self.trainer.lightning_module.log(key, mem, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+            if self.log_gpu_memory == "min_max":
+                self.trainer.lightning_module.log(key, mem, prog_bar=False, logger=True)
+            else:
+                gpu_id = int(key.split("/")[0].split(":")[1])
+                if gpu_id in self.trainer.accelerator_connector.parallel_device_ids:
+                    self.trainer.lightning_module.log(
+                        key, mem, prog_bar=False, logger=True, on_step=True, on_epoch=False
+                    )
 
     """
     Utilities and properties

tests/trainer/logging_/test_train_loop_logging.py

@@ -646,18 +646,22 @@ def training_step(self, batch, batch_idx):
 
 
 @RunIf(min_gpus=2)
-def test_log_gpu_memory_without_logging_on_step(tmpdir):
+@pytest.mark.parametrize("log_gpu_memory", ["all", "min_max"])
+def test_log_gpu_memory_without_logging_on_step(tmpdir, log_gpu_memory):
 
     model = BoringModel()
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_epochs=1,
         limit_train_batches=1,
         limit_val_batches=0,
-        log_gpu_memory="all",
+        log_gpu_memory=log_gpu_memory,
         log_every_n_steps=1,
         gpus=[1],
     )
     trainer.fit(model)
-
-    assert "gpu_id: 1/memory.used (MB)" in trainer.logged_metrics
+    if log_gpu_memory == "min_max":
+        assert "min_gpu_mem" in trainer.logged_metrics
+        assert "max_gpu_mem" in trainer.logged_metrics
+    else:
+        assert "gpu_id: 1/memory.used (MB)" in trainer.logged_metrics