文本和损失

import torch

output =  torch.tensor(
        [[[-2.4210, -2.8106, -1.0321]],
        [[-2.4162,  1.4910,  2.2420]],
        [[-0.9467, -3.4340,  1.5137]],
        [[-1.1759,  2.3805, -0.1223]],
        [[-0.6414, -3.3144,  4.2676]],
        [[-0.9885, 16.1116,  0.5920]]])
#2表示找第3维的最大值 (0表示找第1维的最大值 1表示找第2维的最大值 等等)
#keepdim = True 用[]扩起来 (keepdim = False 不用[]括起来)
predict0 = output.data.max(2, keepdim=True)[0]
predict1 = output.data.max(2, keepdim=True)[1