Sempre é um sentimento um pouco misto escrever um código de princípios básicos bastante bom para fazer algum rearranjo de tensores, apenas para descobrir que o PyTorch tem uma função embutida que faz isso mais rápido. Eu tinha me proposto a pelo menos dar uma olhada na documentação de cada função de torch e tensor, mas se você não sabe para que usaria algo, provavelmente não vai ficar na sua cabeça para quando a oportunidade surgir. Pixel_unshuffle, estou olhando para você.