How can I use median as the reduction function for loss aggregation?

[dempsey-ryan]

I want to log the median loss every epoch, and Lightning is fighting me. I originally did self.log(..., reduce_dx=torch.median) which gave me error lightning_fabric.utilities.exceptions.MisconfigurationException: Only `self.log(..., reduce_fx={min,max,mean,sum})` are supported. If you need a custom reduction, please log a `torchmetrics.Metric` instance instead. Found: <function median at 0x7f53bbd59fc0>

This is the implementation of my Torch Metric:

import torch
from torchmetrics import Metric

class MedianMetric(Metric):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.add_state("median", default=torch.Tensor(0))
    
    def update(self, values):
        self.median = torch.median(self._input_format(values))

    def compute(self):
        return self.median.float()
    
    def __name__(self):
        return "median"

in my LightningModule:

class MyModel(pl.LightningModule):
    def __init__(self, ...):
        self._red = MedianMetric()
    ...
    def training_step(self, batch, batch_idx):
        loss, _ = self.calculate_metrics_and_loss(batch, batch_idx, metric=False)
        self.log("train_loss", loss, on_epoch=True, prog_bar=True, reduce_fx=self._red)
        return loss

Now I get TypeError: MedianMetric.update() takes 2 positional arguments but 3 were given. I'm pretty sure it wants me to put preds and inputs as the arguments to update, but then I'm re-calculating the loss once in training_step and then a second time in MedianMetric.compute. Why is this so difficult? All I want is to use median as the reduction function!