Fix ResultCollection._get_cache with multielement tensors (#9582)

Author: carmocca

CHANGELOG.md

@@ -360,6 +360,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed `BasePredictionWriter` not returning the batch_indices in a non-distributed setting ([#9432](https://github.com/PyTorchLightning/pytorch-lightning/pull/9432))
 
 
+- Fixed check on torchmetrics logged whose `compute()` output is a multielement tensor ([#9582](https://github.com/PyTorchLightning/pytorch-lightning/pull/9582))
+
+
 - Fixed `add_argparse_args` raising `TypeError` when args are typed as `typing.Generic` in Python 3.6 ([#9554](https://github.com/PyTorchLightning/pytorch-lightning/pull/9554))
 
 

pytorch_lightning/trainer/connectors/logger_connector/result.py

@@ -467,7 +467,7 @@ def _get_cache(result_metric: ResultMetric, on_step: bool) -> Optional[torch.Ten
         if on_step and result_metric.meta.on_step:
             cache = result_metric._forward_cache
         elif not on_step and result_metric.meta.on_epoch:
-            if not result_metric._computed:
+            if result_metric._computed is None:
                 # always reduce on epoch end
                 should = result_metric.meta.sync.should
                 result_metric.meta.sync.should = True

tests/core/test_metric_result_integration.py

@@ -27,7 +27,12 @@
 import tests.helpers.utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
-from pytorch_lightning.trainer.connectors.logger_connector.result import ResultCollection
+from pytorch_lightning.trainer.connectors.logger_connector.result import (
+    _Metadata,
+    _Sync,
+    ResultCollection,
+    ResultMetric,
+)
 from pytorch_lightning.utilities.imports import _fault_tolerant_training, _TORCH_GREATER_EQUAL_1_7
 from tests.helpers import BoringModel
 from tests.helpers.runif import RunIf
@@ -544,3 +549,16 @@ def on_train_epoch_end(self) -> None:
 
     trainer = Trainer(default_root_dir=tmpdir, max_epochs=2, limit_train_batches=2, limit_val_batches=0)
     trainer.fit(model)
+
+
+def test_metric_result_computed_check():
+    """Unittest ``_get_cache`` with multielement tensors."""
+    sync = _Sync()
+    metadata = _Metadata("foo", "bar", on_epoch=True, enable_graph=True)
+    metadata.sync = sync
+    rm = ResultMetric(metadata, is_tensor=True)
+    computed_value = torch.tensor([1, 2, 3])
+    rm._computed = computed_value
+    cache = ResultCollection._get_cache(rm, on_step=False)
+    # `enable_graph=True` so no detach, identity works
+    assert cache is computed_value