def augment(self, args):
        """
        Masking the time channel (shape[0])

        Parameters
        ----------
        spectrogram : tf.Tensor, shape [T, num_feature_bins] or [T, num_feature_bins, 1]
            Audio features

        Returns
        -------
        tf.Tensor, shape [T, num_feature_bins] or [T, num_feature_bins, 1]
            Masked time dim of audio features
        """
        with tf.name_scope("time_masking_specaugment"):
            spectrogram, spectrogram_length = args
            max_length, *rest = shape_util.shape_list(spectrogram, out_type=tf.int32)
            indices_shape = (-1,) + (1,) * len(rest)
            mval = get_mask_value(spectrogram, mask_value=self.mask_value)
            T = tf.cast(tf.floor(tf.cast(spectrogram_length, dtype=spectrogram.dtype) * self.p_upperbound), dtype=tf.int32)
            for _ in range(self.num_masks):
                prob = tf.random.uniform(shape=[], minval=0, maxval=1, dtype=spectrogram.dtype)
                do_apply = tf.where(tf.less_equal(prob, self.prob), tf.constant(1, tf.int32), tf.constant(0, tf.int32))
                t = tf.random.uniform(shape=[], minval=0, maxval=T, dtype=tf.int32)
                t = do_apply * tf.minimum(t, spectrogram_length)
                t0 = do_apply * tf.random.uniform(shape=[], minval=0, maxval=(spectrogram_length - t), dtype=tf.int32)
                indices = tf.reshape(tf.range(max_length), indices_shape)
                condition = tf.math.logical_and(tf.math.greater_equal(indices, t0), tf.math.less(indices, t0 + t))
                spectrogram = tf.where(condition, mval, spectrogram)
            return spectrogram, spectrogram_length

        Masking the time channel (shape[0])

        Parameters
        ----------
        spectrogram : tf.Tensor, shape [T, num_feature_bins] or [T, num_feature_bins, 1]
            Audio features

        Returns
        -------
        tf.Tensor, shape [T, num_feature_bins] or [T, num_feature_bins, 1]
            Masked time dim of audio features
        
tensorflow_asr/augmentations/methods/specaugment.py