import tensorflow as tf

    mask_loss = losses.mask_rcnn_loss(
        model_outputs['mask_outputs'], model_outputs['mask_targets'],
        model_outputs['select_class_targets'], params)
  else:
    mask_loss = 0.
  if variable_filter_fn:
    var_list = variable_filter_fn(tf.trainable_variables(),
                                  params['resnet_depth'])
  else:
    var_list = None
  l2_regularization_loss = _WEIGHT_DECAY * tf.add_n(
      [tf.nn.l2_loss(v) for v in var_list
       if 'batch_normalization' not in v.name and 'bias' not in v.name])
  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']:
