Siempre es una sensación un poco contradictoria escribir código de principios básicos bastante bueno para hacer algún reajuste de tensores, solo para descubrir que PyTorch tiene una función incorporada que lo hace más rápido. Me propuse al menos hojear la documentación de cada antorcha y función tensorial, pero si no sabes para qué usarías algo, probablemente no te quedará grabado cuando surja la oportunidad. Pixel_unshuffle, te estoy mirando.