import tensorflow as tf

  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

  inst_weights = gather_init_weights()
  bs = FLAGS.train_batch_size
  hw = FLAGS.src_hw
  inst_weights, indices = tf.nn.top_k(
      inst_weights,
      k=bs,
      sorted=True,
  )
  src_features = tf.reshape(src_features, [
      bs * FLAGS.source_train_batch_multiplier,
      hw,
      hw,
      1,
  ])
  src_features = tf.gather(src_features, indices, axis=0)
  src_features = tf.stop_gradient(src_features)
