对于二维张量，飞话不多说，直接看代码

input_tensor = torch.tensor([[2,3,5],[2,2,2],[3,1,3]], dtype=torch.float32) 
# input_tensor = torch.rand((1000,1000,100))
print(input_tensor)
print("*"*40)
print("沿着dim=1，计算下面两行代码求和，得到每一行权重之和为1")
soft_tensor = torch.softmax(input_tensor, dim=1)
print(soft_tensor, "\n", soft_tensor[0].sum(), soft_tensor[1].sum(), soft_tensor[2].sum())
print("沿着dim=0，计算下面两行代码求和，得到每一行权重之和不为1")
soft_tensor = torch.softmax(input_tensor, dim=0)
print(soft_tensor, "\n", soft_tensor[0].sum(), soft_tensor[1].sum(), soft_tensor[2].sum())

运行结果：

dim维度 dim=0代表是列,dim=1代表是行

由上述运行结果可以发现，沿着dim=1行求softmax，使得每一行求和为1。