img is a (H, W, C) tensor (height, width, channels), the common image layout. Return it rearranged to (C, H, W), the layout most PyTorch models expect.

Input

tensor([[[ 0.,  1.,  2.,  3.],
         [ 4.,  5.,  6.,  7.],
         [ 8.,  9., 10., 11.]],

        [[12., 13., 14., 15.],
         [16., 17., 18., 19.],
         [20., 21., 22., 23.]]])

Output

tensor([[[ 0.,  4.,  8.],
         [12., 16., 20.]],

        [[ 1.,  5.,  9.],
         [13., 17., 21.]],

        [[ 2.,  6., 10.],
         [14., 18., 22.]],

        [[ 3.,  7., 11.],
         [15., 19., 23.]]])