From PyTorch 1.6 to 1.7.0.dev20200914, I am able to export the PyTorch Lightning Module torch script.
But starting from 1.7.0.dev20200915, it started to show error like this.
RuntimeError: 
Wrong type for attribute assignment. Expected None but got Any:

Could this be there are some changes to the JIT? If so, how can I fix it so it can be exported in 1.6 and upcoming 1.7.
Reproducible Link:
https://colab.research.google.com/drive/1DuVwIZ1XHPdyzIUiQcP5jG1O13Zv6K7z?usp=sharing 5
Thank you in advance.