There is probably a really simple way to do this but I’m just not seeing it, I want to obtain the following behavior using torch commands: I have a 2D tensor with dimension (N*n x n) , i.e. a tensor of N , nxn sub-matrices. I want to reshape it from being a stack of the block matrices (as it is now), into a horizontal row of these sub-matrices (dimension (n x N*n)) (i.e)

(A)

(B) —> (A B)

As an example if my tensor is (4, 2)

x=torch.tensor([[1,2],[3,4],[6,7],[8,9]])

tensor([[1, 2],

[3, 4],

[6, 7],

[8, 9]])

then the desired output would be a (2,4) of the following form

tensor([[1, 2, 6, 7],

[3, 4, 8, 9]])

where in this case the block sub-matrices are

A=[[1, 2], [3, 4]]

B=[[6, 7], [8, 9]]

Any help is always appreciated!