I don’t think I can explain it much better than wikipedia’s article on soft max 3, but
torch.nn.functional.softmax(score, dim=1)[:, 1] is actually a smooth approximation of torch.argmax(score, dim=1). That is just what you need and the reason almost any classification works with argmax. We write a bit about that in section 7.2 of our book which can be freely downloaded at the moment 3 (7.2.3-7.2.5 in particular). We also use the dice loss in chapter 13 of the book.
Best regards
Thomas
P.S.: If you enclose your code in lines with triple backticks (```), you get source code formatting.