import tensorflow as tf

            cast_outputs_to_float(v)
          else:
            d[k] = tf.cast(v, tf.float32)
      cast_outputs_to_float(model_outputs)
  else:
    model_outputs = _model_outputs()

  # First check if it is in PREDICT mode.
  if mode == tf.estimator.ModeKeys.PREDICT:
    predictions = {}
    predictions['detections'] = model_outputs['detections']
    predictions['image_info'] = features['image_info']
    if params['include_mask']:
      predictions['mask_outputs'] = tf.nn.sigmoid(model_outputs['mask_outputs'])

    if params['use_tpu']:
      return tf.contrib.tpu.TPUEstimatorSpec(mode=mode, predictions=predictions)
    return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

  # Set up training loss and learning rate.
  global_step = tf.train.get_or_create_global_step()
  learning_rate = learning_rates.step_learning_rate_with_linear_warmup(
      global_step,
      params['init_learning_rate'],
      params['warmup_learning_rate'],
      params['warmup_steps'],
