2D Causal convolution PixelCNN fail to learn

Hi all, I am tying to implement a simple 2d causal convolution based PixelCNN. My PyTorch implementation easily reach a sensible per-sample loss around 90+/- but my Julia’s Flux based solution stuck at 200+ loss per sample. I initially thought the problem must be caused by the raster-scan order and column major mismatch. But it is counter-intuitive to think flipping the image by W and H would cause the autoregressive performance deteriorates drastically suddenly. Scanning an image column wisel...