Torch.argmax - I want more than just one value

[ptrblck]

Double post from here.