def fit(
        self, x: np.ndarray, y: np.ndarray, batch_size: int = 128, nb_epochs: int = 10, verbose: bool = False, **kwargs
    ) -> None:
        """
        Fit the classifier on the training set `(x, y)`.

        :param x: Training data.
        :param y: Labels, one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                  shape (nb_samples,).
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param verbose: Display the training progress bar.
        :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
                       "scheduler" which is an optional function that will be called at the end of every
                       epoch to adjust the learning rate.
        """
        import tensorflow as tf

        if self._train_step is None:  # pragma: no cover
            if self._optimizer is None:  # pragma: no cover
                raise ValueError(
                    "An optimizer `optimizer` or training function `train_step` is required for fitting the "
                    "model, but it has not been defined."
                )

            @tf.function
            def train_step(model, images, labels):
                with tf.GradientTape() as tape:
                    input_size = len(labels)

                    outputs = self.model(images, training=True)
                    outputs = tf.reshape(outputs, [input_size, self.gaussian_samples, self.nb_classes])

                    # Classification loss
                    outputs_softmax = tf.reduce_mean(tf.nn.softmax(outputs, axis=2), axis=1)
                    outputs_log_softmax = tf.math.log(outputs_softmax + 1e-10)
                    indices = tf.stack([np.arange(input_size), labels], axis=1)
                    nll_loss = tf.gather_nd(outputs_log_softmax, indices=indices)
                    classification_loss = -tf.reduce_sum(nll_loss)

                    # Robustness loss
                    beta_outputs = outputs * self.beta
                    beta_outputs_softmax = tf.reduce_mean(tf.nn.softmax(beta_outputs, axis=2), axis=1)
                    top2_score, top2_idx = tf.math.top_k(beta_outputs_softmax, k=2)
                    indices_correct = tf.cast(top2_idx[:, 0], labels.dtype) == labels
                    out = tf.boolean_mask(top2_score, indices_correct)
                    out0, out1 = out[:, 0], out[:, 1]
                    icdf_out1 = tf.math.erfinv(2 * out1 - 1) * np.sqrt(2)
                    icdf_out0 = tf.math.erfinv(2 * out0 - 1) * np.sqrt(2)
                    robustness_loss = icdf_out1 - icdf_out0
                    indices = (
                        ~tf.math.is_nan(robustness_loss)
                        & ~tf.math.is_inf(robustness_loss)
                        & (tf.abs(robustness_loss) <= self.gamma)
                    )
                    out0, out1 = out0[indices], out1[indices]
                    icdf_out1 = tf.math.erfinv(2 * out1 - 1) * np.sqrt(2)
                    icdf_out0 = tf.math.erfinv(2 * out0 - 1) * np.sqrt(2)
                    robustness_loss = icdf_out1 - icdf_out0 + self.gamma
                    robustness_loss = tf.reduce_sum(robustness_loss) * self.scale / 2

                    # Final objective function
                    loss = classification_loss + self.lmbda * robustness_loss
                    loss /= input_size

                gradients = tape.gradient(loss, model.trainable_variables)
                self._optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        else:
            train_step = self._train_step

        scheduler = kwargs.get("scheduler")

        y = check_and_transform_label_format(y, nb_classes=self.nb_classes)

        # Apply preprocessing
        x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=True)

        # Check label shape
        if self._reduce_labels:
            y_preprocessed = np.argmax(y_preprocessed, axis=1)

        train_ds = tf.data.Dataset.from_tensor_slices((x_preprocessed, y_preprocessed)).shuffle(10000).batch(batch_size)

        for epoch in trange(nb_epochs, disable=not verbose):
            for images, labels in train_ds:
                # Tile samples for Gaussian augmentation
                input_size = len(images)
                new_shape = [input_size * self.gaussian_samples]
                new_shape.extend(images[0].shape)
                images = tf.reshape(tf.tile(images, (1, 1, 1, self.gaussian_samples)), new_shape)

                # Add random noise for randomized smoothing
                noise = tf.random.normal(shape=images.shape, mean=0.0, stddev=self.scale)
                noisy_inputs = images + noise

                train_step(self.model, noisy_inputs, labels)

            if scheduler is not None:
                scheduler(epoch)

        Fit the classifier on the training set `(x, y)`.

        :param x: Training data.
        :param y: Labels, one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                  shape (nb_samples,).
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param verbose: Display the training progress bar.
        :param kwargs: Dictionary of framework-specific arguments. This parameter currently only supports
                       "scheduler" which is an optional function that will be called at the end of every
                       epoch to adjust the learning rate.
        
art/estimators/certification/randomized_smoothing/macer/tensorflow.py