Getting issues with torch.onnx.export

I would keep it short. I have created a colab for your easy referencing. Model used is an open source model and I want to create computational graph of it using torch.onnx.export (currently getting Assertion error). I have crafted a dummy_data (this cell) for exporting that can be changed as you wish.