mirror of
https://github.com/twitter/the-algorithm-ml.git
synced 2024-12-23 14:51:49 +01:00
18 lines
287 B
Python
18 lines
287 B
Python
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
|