Full code and trace https://colab.research.google.com/drive/1XBMF3tOQwRLuPrGJib-YkN20_9nmMVT4#scrollTo=Cgi6jl_6R6fj&line=8&uniqifier=1