Data accumulation while computing in GPU

[ravikakaiya]

🐛 Bug

To Reproduce

create a dataloader.
Passing batch wise images to GPU.
Computing SSIM on GPU.
After each batch the GPU memory allocation keeps on increasing.
Resulting into Cuda outof Memory

Expected behavior

RuntimeError: CUDA out of memory. Tried to allocate 210.00 MiB (GPU 0; 15.78 GiB total capacity; 14.39 GiB already allocated; 138.50 MiB free; 14.42 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Environment

TorchMetrics version (and how you installed TM, e.g. conda, pip, build from source): TorchMetrics version 0.8.2
Python & PyTorch Version (e.g., 1.0): python 3.9.12 , pytorch 1.11.0
Any other relevant information such as OS (e.g., Linux):OS : Linux

[SkafteNicki]

Hi @ravikakaiya,
This is warned about when using the metric: https://github.com/Lightning-AI/metrics/blob/203ab6b13cad0219b484f3e47c34b6e7c8831af1/src/torchmetrics/image/ssim.py#L86-L90
We need to store preds and target for computing over all batches.

If you only want to compute the value on the current batch, you could use the functional implementation
https://torchmetrics.readthedocs.io/en/stable/image/structural_similarity.html#functional-interface

Alternatively, you if you are using the modular implementation you can all metric.reset() after calling metric.compute() to reset the internal accumulation buffer.

Closing issue.