fixing state save bug in training.
fixing state save bug in training.

PiperOrigin-RevId: 222171991

tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor,tensorflow/tensor2tensor