import tensorflow as tf

    # nearest neighbor upsampling. The second option may introduce the "blocking
    # effect" but is computationally efficient.
    if model_options.prediction_with_upsampled_logits:
      logits = _resize_bilinear(logits,
                                #tf.shape(images)[1:3],
                                tf.TensorShape([512,512]),
                                scales_to_logits[MERGED_LOGITS_SCOPE].dtype)
      predictions[output] = tf.argmax(logits, 3, output_type=tf.dtypes.int32)
      #predictions[output + PROB_SUFFIX] = tf.nn.softmax(logits)
    else:
      argmax_results = tf.argmax(logits, 3, output_type=tf.dtypes.int32)
      argmax_results = tf.image.resize_nearest_neighbor(
          tf.expand_dims(argmax_results, 3),
          tf.shape(images)[1:3],
          align_corners=True,
          name='resize_prediction')
      predictions[output] = tf.squeeze(argmax_results, 3)
      #predictions[output + PROB_SUFFIX] = tf.image.resize_bilinear(
      #    tf.nn.softmax(logits),
      #    tf.shape(images)[1:3],
      #    align_corners=True,
      #    name='resize_prob')
  return predictions

