How to implement a custom Metric based on sklearn's functions
❓ Questions and Help
What is your question?
How to implement a custom Metric based on sklearn's functions?
Code
I tried to implement a ROCAUC Metric with help of sklearn's roc_auc_score, because it supports multilabel classification.
from pytorch_lightning.metrics import Metric
from sklearn.metrics import roc_auc_score
from typing import Optional, Any
import torch


class ROCAUC(Metric):
    def __init__(
        self,
        average='macro',
        compute_on_step: bool = True,
        dist_sync_on_step: bool = False,
        process_group: Optional[Any] = None,
    ):
        super().__init__(
            compute_on_step=compute_on_step,
            dist_sync_on_step=dist_sync_on_step,
            process_group=process_group,
        )

        self.average = average
        self.add_state("preds", default=[], dist_reduce_fx=None)
        self.add_state("target", default=[], dist_reduce_fx=None)

    def update(self, preds: torch.Tensor, target: torch.Tensor):
        """
        Update state with predictions and targets.

        Args:
            preds: Predictions from model
            target: Ground truth values
        """
        self.preds.append(preds)
        self.target.append(target)

    def compute(self):
        preds = torch.cat(self.preds, dim=0).cpu().numpy()
        target = torch.cat(self.target, dim=0).cpu().numpy().astype(int)
        score = roc_auc_score(target, preds, average=self.average)
        return torch.tensor([score], dtype=torch.float32)
In the training procedure, I log the ROCAUC score along with other metrics at the end of validation epoch.
# this is the code in LightningModule
def validation_epoch_end(self, outputs):
    for name, metric in self.metrics.items():
        value = metric.compute()
        self.log(
            name,
            value,
            prog_bar=True,
            logger=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True)
It worked fine with single GPU training (without ddp).
However, when using ddp accelerator  and setting move_metrics_to_cpu = True . It threw the following error:
Traceback (most recent call last):
  File "train.py", line 75, in <module>
    main()
  File "train.py", line 71, in main
    train_classifier(cfg, train_loader, val_loader, work_dir, seed, **args)
  File "/home/sandylaker/PycharmProjects/xray-classification/xray_cls/apis/train.py", line 40, in train_classifier
    trainer.fit(lit_model, train_loader, val_loader)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 470, in fit
    results = self.accelerator_backend.train()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 152, in train
    results = self.ddp_train(process_idx=self.task_idx, model=model)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 307, in ddp_train
    results = self.train_or_test()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 69, in train_or_test
    results = self.trainer.train()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 492, in train
    self.run_sanity_check(self.get_model())
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 690, in run_sanity_check
    _, eval_results = self.run_evaluation(test_mode=False, max_batches=self.num_sanity_val_batches)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 622, in run_evaluation
    deprecated_eval_results = self.evaluation_loop.evaluation_epoch_end()
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 208, in evaluation_epoch_end
    deprecated_results = self.__run_eval_epoch_end(self.num_dataloaders, using_eval_result)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 246, in __run_eval_epoch_end
    eval_results = model.validation_epoch_end(eval_results)
  File "/home/sandylaker/PycharmProjects/xray-classification/xray_cls/apis/lit_model.py", line 97, in validation_epoch_end
    self.log(
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 266, in log
    self._results.log(
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 146, in log
    value = sync_fn(value, group=sync_dist_group, reduce_op=sync_dist_op)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/accelerators/ddp_accelerator.py", line 344, in sync_tensor
    return sync_ddp_if_available(tensor, group, reduce_op)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py", line 124, in sync_ddp_if_available
    return sync_ddp(result, group=group, reduce_op=reduce_op)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py", line 155, in sync_ddp
    torch.distributed.all_reduce(result, op=op, group=group, async_op=False)
  File "/home/sandylaker/anaconda3/envs/ml/lib/python3.8/site-packages/torch/distributed/distributed_c10d.py", line 954, in all_reduce
    work = _default_pg.allreduce([tensor], opts)
RuntimeError: Tensors must be CUDA and dense
It seemed to be caused by sync_dist=True in the log function. When I modified the last line in ROCAUC's compute function:
return torch.tensor([score], dtype=torch.float32, device=self.preds[0].device)
the error is solved.
When we wrap the sklearn's metric function into pytorch-lightning's Metric, we have to first convert Tensor to numpy.ndarray, then we have to convert numpy.ndarray to Tensor after computing the metric value.
So my question is: Is there a better way than manually moving the Tensor to GPU like this?
What's your environment?

OS: Ubuntu 20.04
Packaging: pip
Version: 1.1.2