def _batch_process(
        self, i_epoch: int, nb_epochs: int, batch_id: int, x_batch: np.ndarray, y_batch: np.ndarray
    ) -> tuple[float, float, float]:
        """
        Perform the operations of OAAT for a batch of data.
        See class documentation for more information on the exact procedure.

        :param i_epoch: training epoch number.
        :param nb_epochs: total training epochs.
        :param batch_id: batch_id of input data.
        :param x_batch: batch of x.
        :param y_batch: batch of y.
        :return: Tuple containing batch data loss, batch data accuracy and number of samples in the batch
        """
        import torch
        from torch import nn
        import torch.nn.functional as F

        if self._classifier.optimizer is None:
            raise ValueError("Optimizer of classifier is currently None, but is required for adversarial training.")

        if self._proxy_classifier.optimizer is None:
            raise ValueError(
                "Optimizer of proxy classifier is currently None, but is required for adversarial training."
            )

        self._classifier.model.train(mode=False)

        if self._apply_lpips_pert and (batch_id % 2 != 0):
            x_batch_pert = self._attack_lpips(
                x_batch,
                y=y_batch,
                eps=self._train_params["i_epsilon"],
                eps_step=self._train_params["i_epsilon"] / 4.0,
                max_iter=self._train_params["max_iter"],
                training_mode=False,
            )

        elif self._apply_lpips_pert and (batch_id % 2 == 0):
            self._attack.set_params(
                eps=self._train_params["mixup_epsilon"],
                eps_step=self._train_params["mixup_epsilon"] / 4.0,
                max_iter=self._train_params["max_iter"],
            )
            x_batch_pert = self._attack.generate(x_batch, y=y_batch)

        else:
            if i_epoch <= (nb_epochs // 4):
                self._attack.set_params(
                    eps=self._train_params["i_epsilon"], eps_step=self._train_params["i_epsilon"] / 2.0, max_iter=5
                )
            else:
                self._attack.set_params(
                    eps=self._train_params["i_epsilon"],
                    eps_step=self._train_params["i_epsilon"] / 4.0,
                    max_iter=self._train_params["max_iter"],
                )

            x_batch_pert = self._attack.generate(x_batch, y=y_batch)

        # Apply preprocessing
        y_batch = check_and_transform_label_format(y_batch, nb_classes=self.classifier.nb_classes)

        x_preprocessed, y_preprocessed = self._classifier._apply_preprocessing(x_batch, y_batch, fit=True)
        x_preprocessed_pert, _ = self._classifier._apply_preprocessing(x_batch_pert, y_batch, fit=True)

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

        i_batch = torch.from_numpy(x_preprocessed).to(self._classifier.device)
        i_batch_pert = torch.from_numpy(x_preprocessed_pert).to(self._classifier.device)
        o_batch = torch.from_numpy(y_preprocessed).to(self._classifier.device)

        self._classifier.model.train(mode=True)

        if self._apply_wp:
            if not self._apply_lpips_pert:
                i_batch_pert_awp = torch.clamp(i_batch + 2.0 * (i_batch_pert - i_batch), 0, 1)
            else:
                i_batch_pert_awp = i_batch_pert
            w_perturb = self._weight_perturbation(x_batch=i_batch, x_batch_pert=i_batch_pert_awp, y_batch=o_batch)
            list_keys = list(w_perturb.keys())
            self._modify_classifier(self._classifier, list_keys, w_perturb, op="add")

        # Zero the parameter gradients
        self._classifier.optimizer.zero_grad()

        n = x_batch.shape[0]
        # Perform prediction
        model_outputs = self._classifier.model(i_batch)
        model_outputs_pert = self._classifier.model(i_batch_pert)

        # Form the loss function
        loss_clean = self._classifier.loss(model_outputs, o_batch)

        if self._apply_lpips_pert and (batch_id % 2 == 0):
            i_batch_pert_eps = torch.min(
                torch.max(i_batch_pert, i_batch - self._train_params["i_epsilon"]),
                i_batch + self._train_params["i_epsilon"],
            )
            model_outputs_pert_eps = self._classifier.model(i_batch_pert_eps)
            loss_kl = (1.0 / n) * nn.KLDivLoss(reduction="sum")(
                F.log_softmax(model_outputs_pert_eps, dim=1),
                torch.clamp(
                    self._train_params["alpha"] * F.softmax(model_outputs, dim=1)
                    + (1.0 - self._train_params["alpha"]) * F.softmax(model_outputs_pert, dim=1),
                    min=EPS,
                ),
            )

        else:
            loss_kl = (1.0 / n) * nn.KLDivLoss(reduction="sum")(
                F.log_softmax(model_outputs_pert, dim=1), torch.clamp(F.softmax(model_outputs, dim=1), min=EPS)
            )
        loss = loss_clean + self._train_params["beta"] * loss_kl

        loss.backward()

        self._classifier.optimizer.step()

        if self._apply_wp:
            self._modify_classifier(self._classifier, list_keys, w_perturb, op="subtract")

        train_loss = loss.item() * o_batch.size(0)
        train_acc = (model_outputs_pert.max(1)[1] == o_batch).sum().item()
        train_n = o_batch.size(0)

        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 i_epoch == epoch_swa:
                for name, params in self._classifier.model.state_dict().items():
                    p_classifier.model.state_dict()[name] = params
            elif i_epoch > epoch_swa:
                for name, params in self._classifier.model.state_dict().items():
                    p_classifier.model.state_dict()[name] = (
                        1 - tau_val
                    ) * params + tau_val * p_classifier.model.state_dict()[name]
            else:
                pass

        self._classifier.model.train(mode=False)

        return train_loss, train_acc, train_n

        Perform the operations of OAAT for a batch of data.
        See class documentation for more information on the exact procedure.

        :param i_epoch: training epoch number.
        :param nb_epochs: total training epochs.
        :param batch_id: batch_id of input data.
        :param x_batch: batch of x.
        :param y_batch: batch of y.
        :return: Tuple containing batch data loss, batch data accuracy and number of samples in the batch
        
art/defences/trainer/adversarial_trainer_oaat_pytorch.py