def fit(
        self,
        x: np.ndarray,
        y: np.ndarray,
        validation_data: tuple[np.ndarray, np.ndarray] | None = None,
        batch_size: int = 128,
        nb_epochs: int = 20,
        **kwargs,
    ):
        """
        Train a model adversarially with OAAT protocol.
        See class documentation for more information on the exact procedure.

        :param x: Training set.
        :param y: Labels for the training set.
        :param validation_data: Tuple consisting of validation data, (x_val, y_val)
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for trainings.
        :param kwargs: Dictionary of framework-specific arguments. These will be passed as such to the `fit` function of
                                  the target classifier.
        """
        import torch

        logger.info("Performing adversarial training with OAAT protocol")

        best_acc_adv_test = 0
        nb_batches = int(np.ceil(len(x) / batch_size))
        ind = np.arange(len(x))

        logger.info("Adversarial Training OAAT")

        if (len(self._train_params["list_swa_epoch"]) != len(self._train_params["list_swa_tau"])) or (
            len(self._train_params["list_swa_epoch"]) != len(self._list_avg_models)
        ):
            raise ValueError(
                "number of elements of list_swa_epoch must be same as of  " "list_swa_tau and list_avg_models."
            )

        beta_init = self._train_params["beta"]
        y = check_and_transform_label_format(y, nb_classes=self.classifier.nb_classes)

        for i_epoch in trange(nb_epochs, desc="Adversarial Training OAAT - Epochs"):

            if i_epoch >= self._train_params["awp_warmup"]:
                self._apply_wp = True

            if i_epoch == (int(3 * nb_epochs / 4)):
                self._classifier._optimizer = torch.optim.SGD(
                    self._classifier.model.parameters(),
                    lr=self._train_params["lr"],
                    momentum=self._train_params["momentum"],
                    weight_decay=self._train_params["weight_decay"],
                )

            if isinstance(self._classifier.optimizer, torch.optim.SGD):

                if self._train_params["oaat_warmup"] == 1:
                    if i_epoch < 10:
                        self.update_learning_rate(self._classifier.optimizer, i_epoch, nb_epochs, lr_schedule="linear")
                    else:
                        self.update_learning_rate(
                            self._classifier.optimizer,
                            i_epoch,
                            nb_epochs,
                            lr_schedule=self._train_params["lr_schedule"],
                        )
                else:

                    self.update_learning_rate(
                        self._classifier.optimizer, i_epoch, nb_epochs, lr_schedule=self._train_params["lr_schedule"]
                    )

            if i_epoch > (nb_epochs // 4):
                self._train_params["i_epsilon"] = i_epoch * self._train_params["epsilon"] / nb_epochs

            if self._train_params["i_epsilon"] > self._train_params["alternate_iter_eps"]:
                self._apply_lpips_pert = True

                load_swa_model_tau = self._train_params["load_swa_model_tau"]
                file_name = f"oaat_swa_tau_{load_swa_model_tau}.model"
                if self._train_params["models_path"] is None:
                    full_path_load = os.path.join(config.ART_DATA_PATH, file_name)
                else:
                    full_path_load = os.path.join(self._train_params["models_path"], file_name)
                if os.path.isfile(full_path_load):
                    self._lpips_classifier._model._model.load_state_dict(torch.load(full_path_load))
                else:
                    raise ValueError("Invalid path/file for weight average model is provided for adversarial training.")

                self._lpips_classifier.model.train(mode=False)
                self._lpips_classifier.model.to(self._classifier.device)

                self._train_params["alpha"] = self._train_params["alpha"] - self._train_params["mixup_alpha"] / (
                    nb_epochs - int(3 * nb_epochs / 4) + 1
                )
                self._train_params["i_lpips_weight"] = (
                    self._train_params["lpips_weight"]
                    * (i_epoch - int(3 * nb_epochs / 4))
                    / (nb_epochs - int(3 * nb_epochs / 4))
                )
                self._train_params["beta"] = beta_init + (self._train_params["beta_final"] - beta_init) * (
                    i_epoch - int(3 * nb_epochs / 4)
                ) / (nb_epochs - int(3 * nb_epochs / 4))

            # Shuffle the examples
            np.random.shuffle(ind)
            start_time = time.time()
            train_loss = 0.0
            train_acc = 0.0
            train_n = 0.0

            for batch_id in range(nb_batches):
                # Create batch data
                x_batch = x[ind[batch_id * batch_size : min((batch_id + 1) * batch_size, x.shape[0])]].copy()
                y_batch = y[ind[batch_id * batch_size : min((batch_id + 1) * batch_size, x.shape[0])]]

                _train_loss, _train_acc, _train_n = self._batch_process(i_epoch, nb_epochs, batch_id, x_batch, y_batch)

                train_loss += _train_loss
                train_acc += _train_acc
                train_n += _train_n

            train_time = time.time()

            # compute accuracy
            if validation_data is not None:
                (x_test, y_test) = validation_data
                y_test = check_and_transform_label_format(y_test, nb_classes=self.classifier.nb_classes)

                x_preprocessed_test, y_preprocessed_test = self._classifier._apply_preprocessing(
                    x_test,
                    y_test,
                    fit=True,
                )
                # pylint: enable=protected-access
                output_clean = np.argmax(self.predict(x_preprocessed_test), axis=1)
                nb_correct_clean = np.sum(output_clean == np.argmax(y_preprocessed_test, axis=1))
                self._attack.set_params(
                    eps=self._train_params["epsilon"],
                    eps_step=self._train_params["epsilon"] / 4.0,
                    max_iter=self._train_params["max_iter"],
                )
                x_test_adv = self._attack.generate(x_preprocessed_test, y=y_preprocessed_test)

                x_preprocessed_test_adv, y_preprocessed_test = self._classifier._apply_preprocessing(
                    x_test_adv,
                    y_test,
                    fit=True,
                )
                # pylint: enable=protected-access
                output_adv = np.argmax(self.predict(x_preprocessed_test_adv), axis=1)
                nb_correct_adv = np.sum(output_adv == np.argmax(y_preprocessed_test, axis=1))

                logger.info(
                    "epoch: %s time(s): %.1f loss: %.4f acc-adv (tr): %.4f acc-clean (val): %.4f acc-adv (val): %.4f",
                    i_epoch,
                    train_time - start_time,
                    train_loss / train_n,
                    train_acc / train_n,
                    nb_correct_clean / x_test.shape[0],
                    nb_correct_adv / x_test.shape[0],
                )

                # save last checkpoint
                if i_epoch + 1 == nb_epochs:
                    self._classifier.save(filename=f"oaat_epoch_{i_epoch}", path=self._train_params["models_path"])

                # save best checkpoint
                if nb_correct_adv / x_test.shape[0] > best_acc_adv_test:
                    self._classifier.save(filename="oaat_epoch_best", path=self._train_params["models_path"])
                    best_acc_adv_test = nb_correct_adv / x_test.shape[0]

            else:
                logger.info(
                    "epoch: %s time(s): %.1f loss: %.4f acc-adv: %.4f",
                    i_epoch,
                    train_time - start_time,
                    train_loss / train_n,
                    train_acc / train_n,
                )

            if i_epoch > self._train_params["swa_save_epoch"] - 1:
                for epoch_swa, tau_val, p_classifier in zip(
                    self._train_params["list_swa_epoch"], self._train_params["list_swa_tau"], self._list_avg_models
                ):
                    if epoch_swa <= i_epoch:
                        file_name = f"oaat_swa_tau_{tau_val}.model"
                        if self._train_params["models_path"] is None:
                            full_path = os.path.join(config.ART_DATA_PATH, file_name)
                        else:
                            full_path = os.path.join(self._train_params["models_path"], file_name)

                        folder = os.path.split(full_path)[0]
                        if not os.path.exists(folder):
                            os.makedirs(folder)

                        # disable pylint because access to _modules required
                        torch.save(p_classifier._model._model.state_dict(), full_path)

        Train a model adversarially with OAAT protocol.
        See class documentation for more information on the exact procedure.

        :param x: Training set.
        :param y: Labels for the training set.
        :param validation_data: Tuple consisting of validation data, (x_val, y_val)
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for trainings.
        :param kwargs: Dictionary of framework-specific arguments. These will be passed as such to the `fit` function of
                                  the target classifier.
        
art/defences/trainer/adversarial_trainer_oaat_pytorch.py