import tensorflow as tf

      logits = tf.log([1.0 - self._relabel_prob, self._relabel_prob])
      mask = tf.squeeze(
          tf.random.categorical(
              logits[None], num_samples=self._sample_batch_size))
