mirror of
https://github.com/twitter/the-algorithm-ml.git
synced 2024-12-23 06:41: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
|