How to instantiate torchmetrics instances properly so Lightning sends them to device

[arzaatri]

The lightning documentation says this about torchmetrics:

"Modular metrics are automatically placed on the correct device when properly defined inside a LightningModule. This means that your data will always be placed on the same device as your metrics. No need to call .to(device) anymore!"

Can anyone expand on what "properly defined" entails? As an example of what may not be proper, I initialized and call my metrics like so:

def get_metrics_for_task(task):
    if task == 'classification:
        task_dict = {
            'accuracy': Accuracy(num_classes=out_size, num_labels=out_size), 
            'precision': Precision(num_classes=out_size, num_labels=out_size), 
            'recall': Recall(num_classes=out_size, num_labels=out_size), 
            'auroc': AUROC(num_classes=out_size, num_labels=out_size), 
        }
    else:
        task_dict = {'mae': MeanAbsoluteError(), 'mse': MeanSquaredError()}
    return task_dict

class MyLightningModule(LightningModule):

    def __init__(...)
        ...
        self.metrics = {k: get_metrics_for_task(task, self.model.out.out_features) for k in ['train', 'val', 'test']}
 
    def validation_step(self, data):
        output, loss = self.run_model(data)
        self.log_metrics(loss, output, data.y, 'val')
        return loss

    def log_metrics(self, loss, output, y phase):
        self.log(f'{phase}_loss', loss)
        for m in self.metrics[phase]:
            metric_obj = self.metrics[phase][m]
            metric_obj(output, y)
            metric_string = f'{phase}_{m}'
            self.log(metric_string, metric_obj, on_step=(phase=='train'), on_epoch=True, batch_size=len(output))

And got this error when running my code:

Encountered different devices in metric calculation (see stacktrace for details). This could be due to the metric class not being on the same device as input. Instead of `metric=BinaryAccuracy(...)` try to do `metric=BinaryAccuracy(...).to(device)` where device corresponds to the device of the input.

Is the way that I define my metrics within the LightningModule improper? If so, what are the requirements to define torchmetrics instance in such a way that Lightning automatically moves them to the correct device?

[Philip-Murzynowski]

https://pytorch-lightning.readthedocs.io/en/1.2.10/extensions/metrics.html#metrics-and-devices:~:text=Being%20properly%20defined%20means