import tensorflow as tf

    ## End new version
    if self._normalize_cols:
      logits_vec = logits_vec - tf.math.reduce_logsumexp(
          logits_vec, axis=0)[None]
    relabel_indices = tf.random.categorical(logits=logits_vec, num_samples=1)

    ### Metrics
    global_step = tf.compat.v1.train.get_or_create_global_step()
    orig_indices = tf.range(
        self._sample_batch_size, dtype=relabel_indices.dtype)
    with tf.name_scope("relabelling"):
      # How often are the originally commanded goals most optimal?
      opt_indices = tf.argmax(logits_vec, axis=1)
      orig_is_opt = opt_indices == orig_indices
