--use_pe doesn't work

Hello kwea! I was studying your Positional Encoding course recently. When running the code, I used the --use_pe and --arch relu parameter, and the generated image looked like it was divided into four parts. There is no such problem when using other parameters. And the psnr of training gradually increases, and the psnr of valid decreases slowly.

I have checked the code for a long time and have not found the problem. I hope to get your help!