Tensorflow tf.nn.depth_to_space vs pixel_shuffle

我想实现一个类似于tf.nn.depth_to_space的函数，但是在Pytorch中。函数pixel_shuffle应该做同样的事情，但我有一个错误

x = tf.constant([[[[1, 2, 3, 15], [4, 5, 6, 16]],
      [[7, 8, 9, 17], [10, 11, 12, 13]]]])
y = tf.nn.depth_to_space(x,2)
output:
array([[[[ 1], [ 2], [ 4],[ 5]],
        [[ 3], [15], [ 6], [16]],
        [[ 7], [ 8], [10], [11]],
        [[ 9], [17], [12], [13]]]], dtype=int32)>

Pytork：

x = torch.Tensor([[[[1, 2, 3, 15], [4, 5, 6, 16]],
      [[7, 8, 9, 17], [10, 11, 12, 13]]]])
pixel_shuffle = torch.nn.PixelShuffle(2)
y = pixel_shuffle(x)

RuntimeError: pixel_shuffle expects its input's 'channel' dimension to be divisible by the square of upscale_factor, but input.size(-3)=2 is not divisible by 4

哪里可能出错，或者我如何在Torch中实现与tf相同的功能

---

谢谢

我认为这个问题是由于Tensorflow和Pytorch的不同数据格式造成的。使用Pytork时，可以将数据转换为NCHW。