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]])