|
import torch
|
|
import torchmetrics as tm
|
|
|
|
import tml.core.metrics as core_metrics
|
|
|
|
|
|
def create_metrics(
|
|
device: torch.device,
|
|
):
|
|
metrics = dict()
|
|
metrics.update(
|
|
{
|
|
"AUC": core_metrics.Auc(128),
|
|
}
|
|
)
|
|
metrics = tm.MetricCollection(metrics).to(device)
|
|
return metrics
|