【扒代码】.max函数 返回最大值&索引

import torch

# 假设输入张量的形状为 [4, 3, 256, 64, 64]
input_tensor = torch.arange(12).reshape(1,3,4)

# 在第 1 维上取最大值
output_tensor,index = input_tensor.max(dim=1)

# print("输入张量形状:", input_tensor.shape)  # 输入张量形状: torch.Size([1, 3, 4])
# print("输出张量形状:", output_tensor.shape)  # 输出张量形状: torch.Size([1, 4])
print(index)    # tensor([[2, 2, 2, 2]])

# 输入张量 : 
# tensor([[[ 0,  1,  2,  3],
#          [ 4,  5,  6,  7],
#          [ 8,  9, 10, 11]]])
# 输出张量 : 
# tensor([[ 8,  9, 10, 11]])