Can you describe in more detail what the problem is? I’m not sure I understand where having a batch size larger than one breaks things, as the max is across the class dimensions and the remaining computations are elementwise or broadcasting across all dimensions.
Out of the box, the code still runs on my end:
>>> import torch
>>> out = torch.randn(8, 3, 100, 100)
>>> mask = out[:, 1, :, :] == torch.max(out, dim=1)[0]
>>> mask = mask.unsqueeze(1)
>>> curr_color = torch.randn(1, 3, 1, 1)
>>> segment = mask*curr_color
>>> segment.shape
torch.Size([8, 3, 100, 100])
>>>