TPU training freeze when define custom backward operation
As is mentioned. A minimal example is here on the colab.
Here is the address for an example. I have noted where it doesn't work. https://colab.research.google.com/gist/rwbfd/57dde430bc168505cfe7d5c42a31924e/tpu_freeze.ipynb