pytorch之stack

a1 = torch.arange(4).view(2,2)
a2 = torch.arange(4,8).view(2,2)
print(a1)
print(a2)
print(torch.stack([a1,a2],dim=0))
print(torch.stack([a1,a2],dim=1))
print(torch.stack([a1,a2],dim=2))
#输出
'''
tensor([[0, 1],
        [2, 3]])
tensor([[4, 5],
        [6, 7]])
tensor([[[0, 1],
         [2, 3]],

        [[4, 5],
         [6, 7]]])
tensor([[[0, 1],
         [4, 5]],

        [[2, 3],
         [6, 7]]])
tensor([[[0, 4],
         [1, 5]],

        [[2, 6],
         [3, 7]]])
'''