run gui_human.py on my own dataset. Problem with visualization

[CHDTevior]

I built my own similar zju-mocap dataset. The size of the image and mask I input is 1088 x 1920, and the input ratio is 0.5, but I reported an error during visualization:

I noticed that the length of the output is 786432=512 x 512 x 3, which happens to be the output of the set zju-mocap dataset 1024 x 1024 x 3 under the corresponding input_ratio. What should I do to change the output to the correct output?
ps: because I found that if forced to use 512 x 512 x 3 as the output of reshape, the visualization will be incomplete:
pred_img = output[f'rgb_level{i}'][b].reshape(512, 512, 3)