Hi,
I just have a follow up question. Since it’s been almost 2 months, here is a refresher.
We wanted to visualize the output of a UNET multi-class segmentation model. The model outputs a tensor of the following form: [batch size, number of classes, h, w] where in each pixel of the tensor, it is an index in [0, number of classes). Obviously if the number of classes isn’t 1 or 3, then this tensor will not be a valid image (also by the pixel values too). So what we do is we map each class index to its corresponding color and then append each class to an output of the form [1, 3, h, w] so that way we can save output as an image in order to visualize it.
for idx, (x, y) in enumerate(loader):
    x = x.to(device=device)
    with torch.no_grad():
        pred = torch.sigmoid(model(x))
        out = model(x)
        class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...]
        output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float)
        for class_idx, color in enumerate(class_to_color):
            mask = out[:,class_idx,:,:] == torch.max(out, dim=1)[0]
            mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
            curr_color = color.reshape(1, 3, 1, 1)
            segment = mask*curr_color # should have shape 1, 3, 100, 100
            output += segment

This solution works perfectly fine for batch sizes of 1. But I tried a batch size of 8 and then mask has the shape [8, 1, 100, 100] and thus something like segment = mask*curr_color won’t work. Do you know how we can modify this so that way it can work for batch sizes that aren’t 1? Thank you!