It takes like 270 seconds per 10 train steps,
while the output example shows 2-5 seconds max.
Is it because we do not have a GPU\TPU support when running on the Coursera
notebook?
I got this warning when create an array using trax.fastmath.numpy:
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)