複雜累加求平均用一行代碼實現，呵呵，我爲什麼要搞這麼無聊的事……
for i in range(0, bs):
sum = 0
dis_max = dis[i][which_max[i]]
sum += dis_max
dis = torch.mean(torch.tensor(sum))
dis_max = torch.mean(torch.tensor(list(dis[i][which_max[i]] for i in range(0, bs)))).cuda()
構建list然後對list進行操作