Refer to following notebook :

https://colab.research.google.com/drive/10Jk9i77Pk8nVuqTnx2HDjwtNfsXzgs8D?usp=sharing