import tensorflow as tf



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

  finetune_one_hot_labels = tf.one_hot(
      tf.cast(labels['finetune'], tf.int64), target_num_classes)
  target_one_hot_labels = tf.one_hot(
      tf.cast(labels['target'], tf.int64), target_num_classes)

  with tf.variable_scope('rl_controller') as rl_scope:
    # It creates a `rl_scope` which will be used for ops.
    pass
  rl_entropy, label_weights, log_prob = rl_label_weights(rl_scope)
  loss_entropy, loss_weights, loss_log_prob = get_loss_weights(rl_scope)

  def gather_init_weights():
    inst_weights = tf.stop_gradient(tf.gather(label_weights, src_labels))
    return inst_weights

