model.to_onnx() fails if self.example_input_array is a list

[t-schanz]

🐛 Bug

Using mode.to_onnx() does not work if the defined example array is a list or tuple of multiple inputs.
The reason is, that it tries to call inputs.to(device) on the list object:

Traceback (most recent call last):
  File "C:\Users\Tobias\Anaconda3\envs\pytorch_local\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 509, in train
    self.train_loop.on_train_end()
  File "C:\Users\Tobias\Anaconda3\envs\pytorch_local\lib\site-packages\pytorch_lightning\trainer\training_loop.py", line 182, in on_train_end
    self.trainer.call_hook('on_train_end')
  File "C:\Users\Tobias\Anaconda3\envs\pytorch_local\lib\site-packages\pytorch_lightning\trainer\trainer.py", line 829, in call_hook
    output = hook_fx(*args, **kwargs)
  File "C:\Users\Tobias\PycharmProjects\QG-GAN\src\utils\QG_GAN.py", line 178, in on_train_end
    self.to_onnx(onnx_file, export_params=True)
  File "C:\Users\Tobias\Anaconda3\envs\pytorch_local\lib\site-packages\pytorch_lightning\core\lightning.py", line 1527, in to_onnx
    input_data = input_data.to(self.device)
AttributeError: 'list' object has no attribute 'to'

Please reproduce using the BoringModel and post here

https://colab.research.google.com/drive/1ZxxIBLRVXF-F0rCn68iYi2WX-GxILS9Z?usp=sharing

Expected behavior

Save the onnx-file, even if the defined self.example_input_array is a list of arrays.

Environment

• CUDA:
  • GPU:
    • GeForce RTX 2080 SUPER
  • available: True
  • version: 10.2
• Packages:
  • numpy: 1.19.1
  • pyTorch_debug: False
  • pyTorch_version: 1.6.0
  • pytorch-lightning: 1.0.3
  • tqdm: 4.49.0
• System:
  • OS: Windows
  • architecture:
    • 64bit
    • WindowsPE
  • processor: Intel64 Family 6 Model 158 Stepping 13, GenuineIntel
  • python: 3.8.5
  • version: 10.0.18362

[rohitgr7]

self.example_input_array = [
    torch.randn([32, 32]),
    torch.randn([32, 32])
]

example_input_array should contain only one tensor sample.

[t-schanz]

@rohitgr7 Thanks for your answer. Could you then please help me to understand how to use that with a forward function that takes multiple inputs, as described in my BoringModel example posted above?

[rohitgr7]

@tmachnitzki here is an example for the same. But yeah, lightning might throw an error here since it expects a Tensor.

It should expect any kind of input IMO as far as it is expected the forward. Will fix this.