import tensorflow as tf

  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)

  src_labels = tf.gather(src_labels, indices)

  inst_weights = bs * inst_weights / tf.reduce_sum(inst_weights)

  src_one_hot_labels = tf.one_hot(tf.cast(src_labels, tf.int64), num_classes)

  src_logits, dst_logits = get_model_logits(src_features, finetune_features,
                                            mode, num_classes,
                                            target_num_classes)

  loss, _, _ = get_final_loss(src_logits, src_one_hot_labels, dst_logits,
