Memory leak when using Metric with list state
🐛 Bug
I tried implementing a custom Metric to use as a loss when training. It seems to compute the desired values fine, however like the title states the metric quickly consumes all the memory on my single GPU. Models that previously required less than half of my GPU memory now run into OOMs after less than one epoch.
I tried replicating the issue in this Colab notebook, however the dataset and training procedure are too lightweight to replicated the resource consumption issue.
Can anyone confirm whether it's an issue with how I handle the list state myself or a bug with Metric itself? The only metric I found that uses a list state that I could draw inspiration from is the explained_variance. However, I thought I should be able to do things differently and wouldn't need to store the targets and preds in the state, only the computed results.
To Reproduce
Colab notebook that tries to reproduce the issue (with no success) on a toy example.
If it can be of further help, my real use-case is the following, where I implement the dice loss as a metric with which to train my model:
class DifferentiableDiceCoefficient(Metric):
    """Computes a differentiable version of the dice coefficient."""

    def __init__(
        self,
        include_background: bool = False,
        nan_score: float = 0.0,
        no_fg_score: float = 0.0,
        reduction: str = "elementwise_mean",
        compute_on_step: bool = True,
        dist_sync_on_step: bool = False,
        process_group: Optional[Any] = None,
    ):
        super().__init__(
            compute_on_step=compute_on_step, dist_sync_on_step=dist_sync_on_step, process_group=process_group
        )
        self.include_background = include_background
        self.nan_score = nan_score
        self.no_fg_score = no_fg_score
        assert reduction in ("elementwise_mean", "none")
        self.reduction = reduction

        self.add_state("dice_by_steps", [])

    def update(self, input: torch.Tensor, target: torch.Tensor) -> None:
        self.dice_by_steps += [
            differentiable_dice_score(
                input=input,
                target=target,
                bg=self.include_background,
                nan_score=self.nan_score,
                no_fg_score=self.no_fg_score,
                reduction=self.reduction,
            )
        ]

    def compute(self) -> torch.Tensor:
        return torch.mean(torch.stack(self.dice_by_steps), 0)
Expected behavior
I would expect to be able to use metrics/losses implemented using the Metric API exactly like I would if they inherited from nn.Module.