How to calculate Confusion matrix?

how can i calculte ROC and confusion matrix for segmentation task

ef get_loss_train(model, data_train, criterion):
    """
        Calculate loss over train set
    """
    model.eval()
    total_acc = 0
    total_loss = 0
    total_dice=0
    result_dice = []
    for batch, (images, masks) in enumerate(data_train):
        with torch.no_grad():
            images = Variable(images.cuda())
            masks = Variable(masks.cuda())
            outputs = model(images).cuda()
            loss = criterion(outputs, masks)
            preds = torch.argmax(outputs, dim=1).float()
            acc = accuracy_check_for_batch(masks.cpu(), preds.cpu(), images.size()[0])
            total_acc = total_acc + acc
            dice = DiceCoefficient()
            score = dice(preds.cpu(), masks.cpu())
            total_dice= score+total_dice
            
            total_loss = total_loss + loss.cpu().item()
    return total_acc/(batch+1), total_loss/(batch + 1),total_dice/(batch+1)

sklearn provides all those metrics

but i do not know i should put them inside the training loop