@vfdev-5 I’m not sure I can use Loss from metrics, no.
So, my update_model function for the trainer engine is
def update_model(engine, batch):
        inputs, targets, weight_map = batch
        inputs = inputs.to(device)
        targets = targets.to(device)
        weight_map = weight_map.to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets, weight_map)
        loss.backward()
        optimizer.step()
        return loss.item()

I need to unpack the batch, that has 3 items, and follow the same scheme for evaluation. If I could unpack the batch inside loss_fn (in metrics.Loss), or using prepare_batch in create_supervised_evaluator that would do the trick.