Optimize Idle time for TPU

[leodzung]

TPU: extreme optimizations

[leodzung]

• The TPU only executes Tensorflow graphs of operations (not Python code)
• the @tf.function annotation is what transforms a piece of Python into a graph.

So the tf.range() loop is running on the TPU itself instead of going back and forth between the TPU and the Kaggle VM.