[feat] Add reset_forward_cache to Metric (#260)

* add reset_forward_cache

* update + tests

* Update torchmetrics/metric.py

Co-authored-by: SkafteNicki <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

CHANGELOG.md

@@ -29,6 +29,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Forward cache is now reset when `reset` method is called ([#260](https://github.com/PyTorchLightning/metrics/pull/260))
 
 ### Deprecated
 

tests/bases/test_metric.py

@@ -303,3 +303,11 @@ def test_warning_on_compute_before_update():
 def test_metric_scripts():
     torch.jit.script(DummyMetric())
     torch.jit.script(DummyMetricSum())
+
+
+def test_metric_forward_cache_reset():
+    metric = DummyMetricSum()
+    _ = metric(2.0)
+    assert metric._forward_cache == 2.0
+    metric.reset()
+    assert metric._forward_cache is None

torchmetrics/metric.py

@@ -166,7 +166,6 @@ def forward(self, *args, **kwargs):
         # add current step
         with torch.no_grad():
             self.update(*args, **kwargs)
-        self._forward_cache = None
 
         if self.compute_on_step:
             self._to_sync = self.dist_sync_on_step
@@ -282,6 +281,7 @@ def reset(self):
         This method automatically resets the metric state variables to their default value.
         """
         self._update_called = False
+        self._forward_cache = None
         # lower lightning versions requires this implicitly to log metric objects correctly in self.log
         if not _LIGHTNING_AVAILABLE or self._LIGHTNING_GREATER_EQUAL_1_3:
             self._computed = None