+1 on open sourcing the code. If you post it somewhere, we'll take a look to see why the Adam optimizer isn't behaving as expected in your implementation.

It may also be helpful to compare your code with the Fashion MNIST Colab example here: https://colab.research.google.com/github/tensorflow/tpu/blob...