How to not throw errors with single-class batches

[cmcmaster1]

Sorry for the very dumb thread, but I still cannot figure out how to only evaluate AUROC at the end of each epoch in PL, so it doesn't throw an error with small batches that do not contain both classes.

[SkafteNicki]

Hi @cmcmaster1,
Here is how you could do your validation loop:

class MyModel(pl.LightningModule):
    def __init__(self, ...):
        super().__init__()
        ...
        self.val_metric = torchmetrics.AUROC()

    def validation_step(self, batch, batch_idx):
        x, y = batch
        yhat = model(x)
        self.metric.update(yhat, y)

    def validation_epoch_end(self, outputs):
        auroc_val = self.metric.compute()
        self.log("val_auroc", auroc_val)

Hope this helps :]

[cmcmaster1]

Helps a lot. Thanks for a very clear and concise snippet of code.