Pytorch объединяет список тензоров

Question

Допустим, у меня есть список тензоров ([A, B, C], где каждый тензор имеет форму [batch_size X 1024].
Я хочу объединить все тензоры в один тензор в следующим образом:
Первая строка в A - это первая строка в новом тензоре, а первая строка в B - вторая строка в новом тензоре, а первая строка в C - третья строка нового тензор и т. д. и т. п.
До сих пор я делал это для циклов, и это неэффективно. Хотелось бы услышать о более эффективных способах.
Спасибо

Answer 1

Вот минимальный пример, который работает:

import torch
a = torch.tensor([[1,1],[1,1]])
b = torch.tensor([[2,2],[2,2]])
c = torch.tensor([[3,3],[3,3]])
torch.stack([a,b,c],dim=0).view(6,2).t().contiguous().view(6,2)

Вывод:

tensor([[1, 1],
        [2, 2],
        [3, 3],
        [1, 1],
        [2, 2],
        [3, 3]])

В вашем случае представление (6,2) должно измениться на batch_size * 3, 1024 .

Решение адаптировано из форумов PyTorch , где показан пример с двумя тензорами.