import tensorflow as tf

  src_loss = get_loss(src_logits, inst_weights, src_one_hot_labels)
  dst_loss = get_loss(dst_logits, 1., finetune_one_hot_labels)
  l2_loss = []
  for v in tf.trainable_variables():
    if 'batch_normalization' not in v.name and 'rl_controller' not in v.name:
      l2_loss.append(tf.nn.l2_loss(v))
  l2_loss = FLAGS.dst_weight_decay * tf.add_n(l2_loss)

  enable_pretrain = tf.cast(
      tf.greater_equal(global_step, FLAGS.first_pretrain_steps), tf.float32)

  loss = src_loss * tf.stop_gradient(loss_weights) * enable_pretrain
  loss += dst_loss + l2_loss

  return tf.identity(loss), src_loss, dst_loss


def train_model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
  """Defines the model function."""
  target_num_classes = FLAGS.target_num_classes
  global_step = tf.train.get_global_step()

  src_features, src_labels = features['src'], tf.cast(labels['src'], tf.int64)
  finetune_features = features['finetune']
  target_features = features['target']

  num_classes = FLAGS.src_num_classes

