import tensorflow as tf

    # There are two ways to obtain the final prediction results: (1) bilinear
    # upsampling the logits followed by argmax, or (2) argmax followed by
    # 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(
