def _weight_perturbation(
        self, x_batch: "torch.Tensor", x_batch_pert: "torch.Tensor", y_batch: "torch.Tensor"
    ) -> dict[str, "torch.Tensor"]:
        """
        Calculate wight perturbation for a batch of data.
        See class documentation for more information on the exact procedure.

        :param x_batch: batch of x.
        :param x_batch_pert: batch of x with perturbations.
        :param y_batch: batch of y.
        :return: dict containing names of classifier model's layers as keys and parameters as values
        """
        import torch
        from torch import nn
        import torch.nn.functional as F

        w_perturb = OrderedDict()
        params_dict, _ = self._calculate_model_params(self._classifier)
        list_keys = list(params_dict.keys())
        self._proxy_classifier.model.load_state_dict(self._classifier.model.state_dict())
        self._proxy_classifier.model.train(mode=True)

        if self._mode.lower() == "pgd":
            # Perform prediction
            model_outputs_pert = self._proxy_classifier.model(x_batch_pert)
            loss = -self._proxy_classifier.loss(model_outputs_pert, y_batch)
        elif self._mode.lower() == "trades":
            n = x_batch.shape[0]
            # Perform prediction
            model_outputs = self._proxy_classifier.model(x_batch)
            model_outputs_pert = self._proxy_classifier.model(x_batch_pert)
            loss_clean = self._proxy_classifier.loss(model_outputs, y_batch)
            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 = -1.0 * (loss_clean + self._beta * loss_kl)

        else:
            raise ValueError(
                "Incorrect mode provided for base adversarial training. 'mode' must be among 'PGD' and 'TRADES'."
            )

        self._proxy_classifier.optimizer.zero_grad()
        loss.backward()
        self._proxy_classifier.optimizer.step()

        params_dict_proxy, _ = self._calculate_model_params(self._proxy_classifier)

        for name in list_keys:
            perturbation = params_dict_proxy[name]["param"] - params_dict[name]["param"]
            perturbation = torch.reshape(perturbation, list(params_dict[name]["size"]))
            scale = params_dict[name]["norm"] / (perturbation.norm() + EPS)
            w_perturb[name] = scale * perturbation

        return w_perturb

        Calculate wight perturbation for a batch of data.
        See class documentation for more information on the exact procedure.

        :param x_batch: batch of x.
        :param x_batch_pert: batch of x with perturbations.
        :param y_batch: batch of y.
        :return: dict containing names of classifier model's layers as keys and parameters as values
        
art/defences/trainer/adversarial_trainer_awp_pytorch.py