From this answer, if your tensor train
has a shape [1000, 19, 1024, 2048]
, you could do :
train_data = train.unfold(2, 64, 64).unfold(3, 64, 64)
.permute(0, 2, 3, 1, 4, 5)
.contiguous()
.view(-1, 19, 64, 64)
It should return the desired tensor of shape [512000, 19, 64, 64]