def fit(
        self,
        x: np.ndarray,
        y: np.ndarray,
        certification_loss: Any = "interval_loss_cce",
        batch_size: int | None = None,
        nb_epochs: int | None = None,
        training_mode: bool = True,
        scheduler: Any | None = None,
        verbose: bool = True,
        **kwargs,
    ) -> None:
        """
        Fit the classifier on the training set `(x, y)`.

        :param x: Training data.
        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                  shape (nb_samples,).
        :param certification_loss: Which certification loss function to use. Either "interval_loss_cce"
                                   or "max_logit_loss". By default, will use interval_loss_cce.
                                   Alternatively, a user can supply their own loss function which takes in as input
                                   the zonotope predictions of the form () and labels of the form () and returns a
                                   scalar loss.
        :param batch_size: Size of batches to use for certified training. NB, this will run the data
                           sequentially accumulating gradients over the batch size.
        :param nb_epochs: Number of epochs to use for training.
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.
        :param scheduler: Learning rate scheduler to run at the start of every epoch.
        :param verbose: If to display the per-batch statistics while training.
        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
               and providing it takes no effect.
        """
        import torch

        if batch_size is None:
            batch_size = self.batch_size

        if nb_epochs is not None:
            epochs: int = nb_epochs
        elif self.batch_size is not None:
            epochs = self.batch_size
        else:
            raise ValueError("Value of `epochs` not defined.")

        # Set model mode
        self.classifier._model.train(mode=training_mode)

        if self.classifier.optimizer is None:  # pragma: no cover
            raise ValueError("An optimizer is needed to train the model, but none is provided.")

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

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

        # Check label shape
        y_preprocessed = self.classifier.reduce_labels(y_preprocessed)

        num_batch = int(np.ceil(len(x_preprocessed) / float(self.pgd_params["batch_size"])))
        ind = np.arange(len(x_preprocessed)).tolist()

        x_cert = np.copy(x_preprocessed)
        y_cert = np.copy(y_preprocessed)

        # Start training
        if self.use_certification_schedule:
            if self.certification_schedule is None:
                certification_schedule_function = DefaultLinearScheduler(
                    step_per_epoch=self.bound / epochs, initial_bound=0.0
                )
        else:
            bound = self.bound

        for _ in tqdm(range(epochs)):
            epoch_non_cert_loss = []
            epoch_non_cert_acc = []

            epoch_cert_loss = []
            epoch_cert_acc = []

            if self.use_certification_schedule:
                bound = certification_schedule_function.step()

            # Shuffle the examples
            random.shuffle(ind)

            # Train for one epoch
            pbar = tqdm(range(num_batch), disable=not verbose)

            for m in pbar:
                certified_loss = torch.tensor(0.0).to(self.classifier.device)
                samples_certified = 0
                # Zero the parameter gradients
                self.classifier.optimizer.zero_grad()

                # get the certified loss
                x_cert, y_cert = shuffle(x_cert, y_cert)
                for i, (sample, label) in enumerate(zip(x_cert, y_cert)):

                    self.set_forward_mode("concrete")
                    concrete_pred = self.classifier.model.forward(np.expand_dims(sample, axis=0))
                    concrete_pred = torch.argmax(concrete_pred)

                    if self.classifier.concrete_to_zonotope is None:
                        if sys.version_info >= (3, 8):
                            eps_bound = np.eye(math.prod(self.classifier.input_shape)) * bound
                        else:
                            eps_bound = np.eye(reduce(lambda x, y: x * y, self.classifier.input_shape)) * bound

                        processed_sample, eps_bound = self.classifier.pre_process(cent=np.copy(sample), eps=eps_bound)
                        processed_sample = np.expand_dims(processed_sample, axis=0)
                    else:
                        processed_sample, eps_bound = self.classifier.concrete_to_zonotope(sample, bound)

                    # Perform prediction
                    self.set_forward_mode("abstract")
                    bias, eps = self.classifier.model.forward(eps=eps_bound, cent=processed_sample)
                    bias = torch.unsqueeze(bias, dim=0)

                    if certification_loss == "max_logit_loss":
                        certified_loss += self.classifier.max_logit_loss(  # type: ignore
                            prediction=torch.cat((bias, eps)),
                            target=torch.from_numpy(np.expand_dims(label, axis=0)).to(self.classifier.device),
                        )
                    elif certification_loss == "interval_loss_cce":
                        certified_loss += self.classifier.interval_loss_cce(
                            prediction=torch.cat((bias, eps)),
                            target=torch.from_numpy(np.expand_dims(label, axis=0)).to(self.classifier.device),
                        )
                    else:
                        certified_loss += certification_loss(torch.cat((bias, eps)), np.expand_dims(label, axis=0))

                    certification_results = []
                    bias = torch.squeeze(bias).detach().cpu().numpy()
                    eps = eps.detach().cpu().numpy()

                    for k in range(self.classifier.nb_classes):
                        if k != concrete_pred:
                            cert_via_sub = self.classifier.certify_via_subtraction(
                                predicted_class=concrete_pred, class_to_consider=k, cent=bias, eps=eps
                            )
                            certification_results.append(cert_via_sub)

                    if all(certification_results) and concrete_pred == label:
                        samples_certified += 1

                    if (i + 1) % batch_size == 0 and i > 0:
                        break

                certified_loss /= batch_size
                epoch_cert_loss.append(certified_loss)
                epoch_cert_acc.append(np.sum(samples_certified) / batch_size)

                # Concrete PGD loss
                i_batch = np.copy(
                    x_preprocessed[ind[m * self.pgd_params["batch_size"] : (m + 1) * self.pgd_params["batch_size"]]]
                ).astype("float32")
                o_batch = y_preprocessed[
                    ind[m * self.pgd_params["batch_size"] : (m + 1) * self.pgd_params["batch_size"]]
                ]

                # Perform prediction
                self.set_forward_mode("concrete")

                if self.augment_with_pgd:
                    self.attack = ProjectedGradientDescent(
                        estimator=self.classifier,
                        eps=self.pgd_params["eps"],
                        eps_step=self.pgd_params["eps_step"],
                        max_iter=self.pgd_params["max_iter"],
                        num_random_init=self.pgd_params["num_random_init"],
                    )
                    i_batch = self.attack.generate(i_batch, y=o_batch)

                self.classifier.model.zero_grad()
                model_outputs = self.classifier.model.forward(i_batch)

                # Form the loss function
                non_cert_loss = self.classifier.concrete_loss(
                    model_outputs, torch.from_numpy(o_batch).to(self.classifier.device)
                )
                epoch_non_cert_loss.append(non_cert_loss)
                epoch_non_cert_acc.append(self.classifier.get_accuracy(model_outputs, o_batch))

                if verbose:
                    pbar.set_description(
                        f"Bound {bound:.3f}: "
                        f"Loss {torch.mean(torch.stack(epoch_non_cert_loss)):.2f} "
                        f"Cert Loss {torch.mean(torch.stack(epoch_cert_loss)):.2f} "
                        f"Acc {np.mean(epoch_non_cert_acc):.2f} Cert Acc {np.mean(epoch_cert_acc):.2f}"
                    )

                loss = certified_loss * self.loss_weighting + non_cert_loss * (1 - self.loss_weighting)
                # Do training
                if self.classifier._use_amp:  # pragma: no cover
                    from apex import amp

                    with amp.scale_loss(loss, self.classifier.optimizer) as scaled_loss:
                        scaled_loss.backward()

                else:
                    loss.backward()

                self.classifier.optimizer.step()

            if scheduler is not None:
                scheduler.step()

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

        :param x: Training data.
        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                  shape (nb_samples,).
        :param certification_loss: Which certification loss function to use. Either "interval_loss_cce"
                                   or "max_logit_loss". By default, will use interval_loss_cce.
                                   Alternatively, a user can supply their own loss function which takes in as input
                                   the zonotope predictions of the form () and labels of the form () and returns a
                                   scalar loss.
        :param batch_size: Size of batches to use for certified training. NB, this will run the data
                           sequentially accumulating gradients over the batch size.
        :param nb_epochs: Number of epochs to use for training.
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.
        :param scheduler: Learning rate scheduler to run at the start of every epoch.
        :param verbose: If to display the per-batch statistics while training.
        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
               and providing it takes no effect.
        
art/defences/trainer/certified_adversarial_trainer_pytorch.py