Numpy array slicing/reshape/concatination
concatenate

There is no need to do the concatenation pairwise np.concatenate([img[:, :, :, i][0] for i in range(img.shape[-1])]) should improve the speed already

numpy.moveaxis

You can use numpy.moveaxis

new_array2 = np.concatenate(np.moveaxis(img[0],(0,1,2), (1,2,0)))


To check whether the result is the same:

assert np.array_equal(new_array, new_array2)


Other improvements

Since you don't use the first axis of the image, you can do np_img = image.data.numpy()[0] so prevent you repeating the [0] in all subsequent code