Solved by ptrblck in post #2 
              
              
                This should work: 
x = torch.tensor([[[[ 0.4008, -0.6662],
                    [-0.4133,  1.3639]],

                   [[-0.8354,  0.6317],
                    [ 0.3240, -1.1438]],

                   [[-0.3452,  1.2110],
                    [ 0.6575,  0.9924]]]])

pred = torch.argmax(x, dim=1)
pri…