If you'd like to gather more information on the error and if the error occurs in the first few iterations, I suggest you run the experiment in CPU-only mode (no GPUs). The error message will be much more specific.
Source: https://github.com/tensorflow/tensor2tensor/issues/574