[mrwilliamchang]

To see memory consumption on the TPU while running on GKE you can look at kubernetes.io/node/accelerator/memory_used

https://cloud.google.com/kubernetes-engine/docs/how-to/tpus#...