Yes the shape is correct.

Here is the code:

```
pred = torch.sigmoid(model(x))
out = (pred > 0.5).float()
print(f"out shape: {out.shape}\n")
class_to_color = [torch.tensor([0.0, 0.0, 0.0]), torch.tensor([10, 133, 1]), torch.tensor([14, 1, 133]), torch.tensor([33, 255, 1]), torch.tensor([243, 5, 247]), torch.tensor([(255, 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]
print(f"{mask}\n")
mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
print(f"mask shape: {mask.shape}\n")
curr_color = color.reshape(1, 3, 1, 1)
print(f"color shape: {color.shape}\n")
segment = mask*color # should have shape 1, 3, 100, 100
output += segment
torchvision.utils.save_image(output, f"{folder}/pred_{idx}.png")
```

Here is the corresponding output:

```
out shape: torch.Size([1, 6, 100, 100])
tensor([[[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
...,
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True]]], device='cuda:0')
mask shape: torch.Size([1, 1, 100, 100])
color shape: torch.Size([3])
```

And the error:

```
segment = mask*color # should have shape 1, 3, 100, 100
RuntimeError: The size of tensor a (100) must match the size of tensor b (3) at non-singleton dimension 3
```