import tensorflow as tf

        })
    return model_outputs

  if params['use_bfloat16']:
    with tf.contrib.tpu.bfloat16_scope():
      model_outputs = _model_outputs()
      def cast_outputs_to_float(d):
        for k, v in sorted(six.iteritems(d)):
