import tensorflow as tf

  total_loss = (total_rpn_loss + total_fast_rcnn_loss + mask_loss +
                l2_regularization_loss)

  host_call = None
  if mode == tf.estimator.ModeKeys.TRAIN:
    optimizer = create_optimizer(learning_rate, params)
    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

    if not params['resnet_checkpoint']:
      scaffold_fn = None
    else:

