Squeeze tensor while logging (#14489)

Co-authored-by: Adrian Wälchli <[email protected]>

Author: 2 people

src/pytorch_lightning/core/module.py

@@ -419,8 +419,7 @@ def log(
                 " but it should not contain information about `dataloader_idx`"
             )
 
-        value = apply_to_collection(value, numbers.Number, self.__to_tensor)
-        apply_to_collection(value, torch.Tensor, self.__check_numel_1, name)
+        value = apply_to_collection(value, (torch.Tensor, numbers.Number), self.__to_tensor, name)
 
         if self.trainer._logger_connector.should_reset_tensors(self._current_fx_name):
             # if we started a new epoch (running its first batch) the hook name has changed
@@ -552,16 +551,15 @@ def __check_not_nested(value: dict, name: str) -> None:
     def __check_allowed(v: Any, name: str, value: Any) -> None:
         raise ValueError(f"`self.log({name}, {value})` was called, but `{type(v).__name__}` values cannot be logged")
 
-    def __to_tensor(self, value: numbers.Number) -> Tensor:
-        return torch.tensor(value, device=self.device)
-
-    @staticmethod
-    def __check_numel_1(value: Tensor, name: str) -> None:
+    def __to_tensor(self, value: Union[torch.Tensor, numbers.Number], name: str) -> Tensor:
+        value = torch.tensor(value, device=self.device)
         if not torch.numel(value) == 1:
             raise ValueError(
                 f"`self.log({name}, {value})` was called, but the tensor must have a single element."
                 f" You can try doing `self.log({name}, {value}.mean())`"
             )
+        value = value.squeeze()
+        return value
 
     def log_grad_norm(self, grad_norm_dict: Dict[str, float]) -> None:
         """Override this method to change the default behaviour of ``log_grad_norm``.

tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py

@@ -28,7 +28,8 @@
 from pytorch_lightning import callbacks, Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, TQDMProgressBar
 from pytorch_lightning.core.module import LightningModule
-from pytorch_lightning.demos.boring_classes import BoringModel, RandomDataset
+from pytorch_lightning.demos.boring_classes import BoringModel, RandomDataset, RandomDictDataset
+from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests_pytorch.helpers.datasets import RandomDictDataset
 from tests_pytorch.helpers.runif import RunIf
@@ -837,3 +838,13 @@ def on_train_start(self):
 
     assert mock_log_metrics.mock_calls == [call(metrics={"foo": 123.0, "epoch": 0}, step=0)]
     assert trainer.max_epochs > 1
+
+
+def test_unsqueezed_tensor_logging():
+    model = BoringModel()
+    trainer = Trainer()
+    trainer.state.stage = RunningStage.TRAINING
+    model._current_fx_name = "training_step"
+    model.trainer = trainer
+    model.log("foo", torch.Tensor([1.2]))
+    assert trainer.callback_metrics["foo"].ndim == 0