Hi, Good job on the repo! So is this code working on tf1 or tf2? Which one do you recommend? (I am running it on colab)