Memory Leak in AveragePrecision

[Rodrigo-P]

🐛 Bug

The gpu memory usage keeps growing when using AveragePrecision. This doesn't happen with other metrics (Accuracy and Precision).

Is this the expected behaviour for this metric?

To Reproduce

Code sample

from torchmetrics import Accuracy, Precision, AveragePrecision
import torch


def get_memory():
    return torch.cuda.memory_allocated(0) / (1024**2)


def loop_metric(metric, repetitions=20, loop_length=100):
    print(type(metric).__name__)
    print(f"{get_memory():.3f}MB", end="")
    for i in range(repetitions):
        for j in range(loop_length):
            metric(pred, target)
        print(f" -> {get_memory():.3f}MB", end="")
    print("\n")


N_LABELS = 16
avg_p = AveragePrecision(task="multilabel", num_labels=N_LABELS).to("cuda")
prec = Precision(task="multilabel", num_labels=N_LABELS).to("cuda")
acc = Accuracy(task="multilabel", num_labels=N_LABELS).to("cuda")

BATCH_SIZE = 2056
target = torch.randint(size=(BATCH_SIZE, N_LABELS), low=0, high=2, device="cuda")
pred = torch.rand(size=(BATCH_SIZE, N_LABELS), device="cuda")

print()
loop_metric(acc)
loop_metric(prec)
loop_metric(avg_p)

Environment

• TorchMetrics: 1.8.2
• Python: 3.13.7
• PyTorch: 2.9.0+cu126
• Happens in Linux and Windows 10

Additional context

I ran into a OOM error when using AveragePrecision during model training.
Using gc.collect() and torch.cuda.empty_cache() wasn't able to prevent it, and I tracked the issue down to this metric.

Is there some function that needs to be called to free this memory?

[justusschock]

yes it is. The memory would be freed up on calls to reset. The issue is that the state management of the metrics' internals is different where in e.g. accuracy you're onnly adding to other numbers, in AveragePrecision you need to save all the states for accurate computation of the result

(Note that without a call to reset, your results might also be wrong as they accumulate potentially across train/val/test or multiple epochs)

[Rodrigo-P]

Calling reset does clear the memory. This issue mentions that self.log() should reset the memory when called, but this doesn't seem to happen in my code.

I can manually reset the metric so OOM is not a problem anymore, but is there a better way to do this?

[justusschock]

self.log resets the memory at the end of the epoch. What's your training setup? in the example above you don't use self.log at all