If you don't mind, with which method did you train the style onnx models that you run with this code? Because when I try with a few train methods on github, it usually outputs like 1x512x32x32 and when I run it with your code, I get 32x32 images. But it works great with your onnx models. I tried to export the output in your onnx with the dynamic_axes parameter with Torch, but this time the codes give an error.
(name: input
type: float32[input_dynamic_axes_1,3,input_dynamic_axes_2,input_dynamic_axes_3]
name: output
type: float32[output_dynamic_axes_1,3,output_dynamic_axes_2,output_dynamic_axes_3]
)