Fix for modalities in the problem class.
Fix for modalities in the problem class.

PiperOrigin-RevId: 213917724

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