Suppose I have tensor with size [2, 64, 64], I want to stack it either horizontally or vertically along the first dimension, so I want to stack [0,:,:] and [1,:,:] and get a [128,64] or a [64,128] tensor.

Clearly this is possible with looping over the first dimension of tensor, but I’m looking for a vectorized solution instead of for-loop solution, any idea?

And also another general quesion, asking for advice honestly! , I really don’t know how can I improve my coding in a vectorized manner, I mean I usually use for-loops as a solution but it seems using for-loop is not recommended in terms of performance, how can I learn to code in a vectorized manner?

I tried to use concat or vstack but didn’t succeed.