TorchMetrics Not working in google colab

[bibinwils]

I have taken the torchmetrics for regression task in google colab. and it's taking so much time and GPU memory to run. So I removed them from the train and validation step and put them only in the test_step. But after the training, in test state also it's not completing the run, and it's keeping on rolling. What should I do?

[SkafteNicki]

Hi, could you please provide a link to your colab notebook or provide some code example for better understanding your problem.

[bibinwils]

Sorry, since I couldn't get any results I stopped in between. when I looked to get the results, still it's not possible to run when torchmetrics are enabled. This is the code: https://colab.research.google.com/drive/12-wApssxCngkWF9u1Rtcug556CGmtBUT3o?usp=sharing can you please look into it?

[Borda]

@bibinwils — this was likely caused by version incompatibility in Colab circa 2021. A lot has improved since. Here's a clean setup that works today:

Install:

!pip install -q lightning torchmetrics

Common causes of GPU memory / slowness with metrics:

1. Calling metric(preds, target) every step — .forward() triggers compute() each time. For metrics that store all predictions (AUROC, PRC), this is expensive. Use .update() + epoch-level .compute():

def training_step(self, batch, batch_idx):
    self.train_metrics.update(y_hat, y)  # cheap
    return loss

def on_train_epoch_end(self):
    self.log_dict(self.train_metrics.compute())  # compute once
    self.train_metrics.reset()

2. Not calling .reset() — state accumulates indefinitely, eating GPU memory.

3. Using non-binned metrics on large datasets — thresholds=None stores ALL predictions. Use thresholds=200 for constant memory:

from torchmetrics.classification import BinaryAUROC
auroc = BinaryAUROC(thresholds=200)  # O(200) memory instead of O(N)

4. Using .clone() for stage separation:

metrics = MetricCollection({...})
self.train_metrics = metrics.clone(prefix="train/")
self.val_metrics = metrics.clone(prefix="val/")

If you're still hitting issues on the current version, a minimal reproducible example would help debug further.

Docs: Overview