How can I interleave 5 PyTorch tensors?

Viewed 841

I have 5 tensors of shape torch.Size([7, 20, 180])

I want to interleave them, one after the other along dim=1. So that my final shape will be torch.Size([7, 100, 180]).

Basically, I want the first element from the first tensor, then the first element from the second tensor, and so on.

1 Answers

If I understood correctly,

import torch
stacked = torch.stack(list_of_tensors, dim=2)
interleaved = torch.flatten(stacked, start_dim=1, end_dim=2)

interleaved is what you need apparently

(tested with pytorch 1.1.0)

Related