import os

import torch
import torch.distributed as dist


def score_function(engine):
    val_metric = engine.state.metrics["val_mean_dice"]
    if dist.is_initialized():
        device = torch.device("cuda:" + os.environ["LOCAL_RANK"])
        val_metric = torch.tensor([val_metric]).to(device)
        dist.all_reduce(val_metric, op=dist.ReduceOp.SUM)
        val_metric /= dist.get_world_size()
        return val_metric.item()
    return val_metric