Davidh

I'm also having the exact same problem, my config is: `trax 1.3.8 mesh-tensorflow 0.1.19 tensor2tensor 1.15.7 tensorboard 2.6.0 tensorboard-data-server 0.6.1 tensorboard-plugin-wit 1.8.0 tensorflow 2.5.0 tensorflow-addons 0.14.0 tensorflow-datasets 4.4.0 tensorflow-estimator 2.5.0...