Hi.
I am trying to convert the FastPhotoStyle model (PyTorch) by NVIDIA 6. The stylization module takes four inputs i.e. the style image, the content image, and two segmentation masks.
I am running the conversion step like so -
torch.onnx.export(stylization_module, [cont_img, styl_img, cont_img, styl_img], 'stylization_module.onnx', input_names=['content', 'style', 'content_seg', 'style_seg'], output_names=['stylized_image'])`

It is producing:
RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type NoneType

Here’s the Colab Notebook 13 for full reproducibility.
Any pointers to resolve the issue would be great.