I suppose you are trying to call transpose on the numpy array not on your tensor right?
Assuming this, you have to give all 4 dimensions to the transpose call, since your array has 4 dimensions.
This works:
np.random.randn(1, 50, 22, 22).transpose(1, 2, 0, 3)

while this throws your error:
np.random.randn(1, 50, 22, 22).transpose(1, 2, 0)

What do the dimensions stand for in your example, [channel, batch_size, width, height]?
EDIT: In the STN example, they are calling make_grid on the images, which returns one image ([channel, w, h]) containing a grid of the “subimages”. Could you try to use it, too?