DeepAR ONNX Export error

• PyTorch-Forecasting version: 0.9.0
• PyTorch version: 1.9
• Python version: 3.8
• Operating System: Linux

Expected behavior

I executed code https://github.com/jdb78/pytorch-forecasting/blob/master/examples/ar.py

Then I tried to export the model to ONNX:
filepath = "model.onnx" input_sample = torch.randn((64,20)) deepar.to_onnx(filepath, input_sample, export_params=True)

I've got the following error:

----> 1 deepar.to_onnx(filepath, input_sample, export_params=True)

~/anaconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
26 def decorate_context(*args, **kwargs):
27 with self.class():
---> 28 return func(*args, **kwargs)
29 return cast(F, decorate_context)
30

~/anaconda3/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py in to_onnx(self, file_path, input_sample, **kwargs)
1893 if "example_outputs" not in kwargs:
1894 self.eval()
-> 1895 kwargs["example_outputs"] = self(input_sample)
1896
1897 torch.onnx.export(self, input_sample, file_path, **kwargs)

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1050 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051 return forward_call(*input, **kwargs)
1052 # Do not call functions when jit is used
1053 full_backward_hooks, non_full_backward_hooks = [], []

~/anaconda3/lib/python3.8/site-packages/pytorch_forecasting/models/deepar/init.py in forward(self, x, n_samples)
306 Forward network
307 """
--> 308 hidden_state = self.encode(x)
309 # decode
310 input_vector = self.construct_input_vector(

~/anaconda3/lib/python3.8/site-packages/pytorch_forecasting/models/deepar/init.py in encode(self, x)
228 """
229 # encode using rnn
--> 230 assert x["encoder_lengths"].min() > 0
231 encoder_lengths = x["encoder_lengths"] - 1
232 input_vector = self.construct_input_vector(x["encoder_cat"], x["encoder_cont"])

IndexError: too many indices for tensor of dimension 2

Could you please guide me on how to export the model to ONNX and how to get the size of the input tensor of the model?
Thank you!

Start with examples/ar.py and the DeepAR forward and encode frames in pytorch_forecasting/models/deepar/init.py shown by the traceback. Check how the model input is structured before to_onnx receives it, then verify that an appropriate sample can export successfully and that the required input structure is documented.