Hey one more doubt,
Based on this softmax example.
def check_accuracy(loader, model, num_labels, time=0, device=DEVICE):
    num_correct = 0
    num_pixels = 0
    dice_score = 0
    
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            preds = model(x)
            y = y.to(device)

            lossfn = nn.CrossEntropyLoss()
            celoss = lossfn(preds, y.long())

            #y = y.unsqueeze(1)

            preds_labels = torch.argmax(preds, dim=1)
    
            num_correct += (preds_labels == y).sum()
            num_pixels += torch.numel(preds_labels)
     
            dice_score = dice_loss(y, preds_labels, num_labels)

    acc = evaluate_segmentation(preds_labels, y, 4, score_averaging='weighted')
    acc.append(dice_score)
    acc.append(celoss.cpu().numpy())

    print_and_save_results(num_correct, num_pixels, acc, time)  

    model.train()

def save_predictions_as_imgs(loader, model, epoch, folder="data/predictions/", device=DEVICE):
    print("=> Saving predictions as images")
    model.eval()
    with torch.no_grad():    
        for idx, (x, _) in enumerate(loader):
            x = x.to(device)
            preds = model(x)
            preds_labels = torch.argmax(preds, dim=1)
           # preds = torch.log_softmax(model(x), 1)
           # preds_labels = torch.argmax(preds, 1)
            
            preds_labels = label_to_pixel(preds_labels)

            now = datetime.now().strftime("%Y%m%d_%H%M%S")

            save_image(preds_labels, f"{folder}{now}_pred_e{epoch}_i{idx}.png")

            # preds_labels = ToPILImage()(preds_labels).convert("L")
            # preds_labels.save(f"{folder}{now}_pred_e{epoch}_i{idx}.png")
            
    model.train()


def save_validation_as_imgs(loader, folder="data/predictions/", device=DEVICE):
    print("=> Saving predictions as images")

    for idx, (_, y) in enumerate(loader):
        y = y.to(device)
        val = y.unsqueeze(1)

        now = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        save_image(val, f"{folder}{now}_val_i{idx}.png")
        #val = ToPILImage()(y).convert("L")
        #val.save(f"{folder}{now}_val_i{idx}.png")


Would this code be correct?