Hi,
@Mellow
This line worked for me.

_, pred = torch.max(scores, dim=1)

https://pytorch.org/docs/stable/torch.html#torch.max
Where scores is your ultimate tensor containing probs.
For instance, scores tensor has size of [10, 150, 256, 256] which means I have 150 classes to segment and using above code give me [10, 256, 256] tensor. Then you can .unsqueeze(1) to get your desired dimension.
PS: torch.max return a tuple which the second value is the target tensor.