在某个索引处连接火炬张量

Posted 2023-03-29

【中文标题】在某个索引处连接火炬张量

【英文标题】：Concatenate torch tensor at a certain index

【发布时间】：2021-07-08 19:54:05

【问题描述】：

我希望在某个索引处连接 2 个火炬张量。例如，我想在 a[1] 之后添加 b。

a = torch.Tensor([1, 2, 3, 4, 5])
b = torch.Tensor([6, 7, 8, 9, 10])

想要的输出是

torch.Tensor([1, 2, 6, 7, 8, 9, 10, 3, 4, 5])

我试过torch.cat，但我只能拥有

tensor([ 6.,  7.,  8.,  9., 10.,  1.,  2.,  3.,  4.,  5.])
tensor([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.])

【参考方案1】：

您需要拆分第一个张量并在其间连接第二个张量

torch.cat([a[:2], b, a[2:]])

输出会是这样的

tensor([ 1.,  2.,  6.,  7.,  8.,  9., 10.,  3.,  4.,  5.])