Fix Metric.state_dict (#5614)

* Fix Metric.state_dict

* Update CHANGELOG.md

* Update CHANGELOG.md

* Detach tensors in a list if needed

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: manipopopo

CHANGELOG.md

@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an error when logging a progress bar metric with a reserved name ([#5620](https://github.com/PyTorchLightning/pytorch-lightning/pull/5620))
 
 
+- Fixed `Metric`'s `state_dict` not included when child modules ([#5614](https://github.com/PyTorchLightning/pytorch-lightning/pull/5614))
 
 
 ## [1.1.5] - 2021-01-19

pytorch_lightning/metrics/metric.py

@@ -284,11 +284,23 @@ def persistent(self, mode: bool = False):
         for key in self._persistent.keys():
             self._persistent[key] = mode
 
-    def state_dict(self, *args, **kwargs):
+    def state_dict(self, destination=None, prefix='', keep_vars=False):
+        destination = super().state_dict(
+            destination=destination,
+            prefix=prefix,
+            keep_vars=keep_vars
+        )
         # Register metric states to be part of the state_dict
-        state_dict = super().state_dict()
         for key in self._defaults.keys():
             if self._persistent[key]:
                 current_val = getattr(self, key)
-                state_dict.update({key: current_val})
-        return state_dict
+                if not keep_vars:
+                    if torch.is_tensor(current_val):
+                        current_val = current_val.detach()
+                    elif isinstance(current_val, list):
+                        current_val = [
+                            cur_v.detach() if torch.is_tensor(cur_v) else cur_v
+                            for cur_v in current_val
+                        ]
+                destination[prefix + key] = current_val
+        return destination

tests/metrics/test_metric.py

@@ -6,6 +6,7 @@
 import numpy as np
 import pytest
 import torch
+from torch import nn
 
 from pytorch_lightning.metrics.metric import Metric
 
@@ -201,3 +202,22 @@ def test_state_dict(tmpdir):
     assert metric.state_dict() == OrderedDict(x=0)
     metric.persistent(False)
     assert metric.state_dict() == OrderedDict()
+
+
+def test_child_metric_state_dict():
+    """ test that child metric states will be added to parent state dict """
+    class TestModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.metric = Dummy()
+            self.metric.add_state('a', torch.tensor(0), persistent=True)
+            self.metric.add_state('b', [], persistent=True)
+            self.metric.register_buffer('c', torch.tensor(0))
+
+    module = TestModule()
+    expected_state_dict = {
+        'metric.a': torch.tensor(0),
+        'metric.b': [],
+        'metric.c': torch.tensor(0)
+    }
+    assert module.state_dict() == expected_state_dict