在某個索引處連接火炬張量

Question

我正在尋找以某個索引連接 2 個火炬張量。 例如，我想在 a[1] 之后添加 b。

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

所需的 output 是

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.])

Answer 1

您將需要拆分第一個張量並在兩者之間連接第二個

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

output 會像

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