如何将两个torch张量交叉拼接

1.问题

之前在编程时遇到一个问题，是将两个a×b大小的张量交叉拼接在一起，在进行后续的处理。比如

[[1,1,1],[1,1,1]]与[[2,2,2],[2,2,2]]交叉拼接在一起的结果为[[1,2,1,2,1,2],[1,2,1,2,1,2]]

2.解决方法

使用torch.stack命令即可。

import torch
aa = torch.ones(2,3)
bb = aa + 1
aa = aa.view(-1,1)
bb = bb.view(-1,1)

cc = torch.stack((aa,bb),dim=1).view(-1,6)
print(cc)

## tensor([[1., 2., 1., 2., 1., 2.],
##        [1., 2., 1., 2., 1., 2.]])