def custom_loss_gradient(
        self,
        loss_fn,
        x: np.ndarray | "torch.Tensor",
        y: np.ndarray | "torch.Tensor",
        layer_name,
        training_mode: bool = False,
    ) -> np.ndarray | "torch.Tensor":
        """
        Compute the gradient of the loss function w.r.t. `x`.

        :param loss_fn: Loss function w.r.t to which gradient needs to be calculated.
        :param x: Sample input with shape as expected by the model(base image).
        :param y: Sample input with shape as expected by the model(target image).
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.`
        :param layer_name: Name of the layer from which activation needs to be extracted/activation layer.
        :return: Array of gradients of the same shape as `x`.
        """
        import torch

        self._model.train(mode=training_mode)
        self._model.eval()
        if self.all_framework_preprocessing:
            if isinstance(x, torch.Tensor):
                x_grad = x.clone().detach().requires_grad_(True)
            else:
                x_grad = torch.tensor(x).to(self._device)
                x_grad.requires_grad = True
            if isinstance(y, torch.Tensor):
                y_grad = y.clone().detach()
            else:
                y_grad = torch.tensor(y).to(self._device)
            inputs_t, _ = self._apply_preprocessing(x_grad, y=None, fit=False, no_grad=False)
            targets_t, _ = self._apply_preprocessing(y_grad, y=None, fit=False, no_grad=False)
        if isinstance(x, np.ndarray):
            x_preprocessed, _ = self._apply_preprocessing(x, y=None, fit=False, no_grad=True)
            y_preprocessed, _ = self._apply_preprocessing(y, y=None, fit=False, no_grad=True)
            x_grad = torch.from_numpy(x_preprocessed).to(self._device)
            y_grad = torch.from_numpy(y_preprocessed).to(self._device)
            x_grad.requires_grad = True
            y_grad.requires_grad = False
            inputs_t = x_grad
            targets_t = y_grad
        else:
            raise NotImplementedError("Combination of inputs and preprocessing not supported.")

        # Compute the gradient and return
        model_outputs1 = self.get_activations(inputs_t, layer_name, 1, framework=True)
        model_outputs2 = self.get_activations(targets_t, layer_name, 1, framework=True)
        diff = model_outputs1 - model_outputs2
        loss = loss_fn(diff, p=2)

        # Clean gradients
        self._model.zero_grad()

        # Compute gradients
        if self._use_amp:  # pragma: no cover
            from apex import amp

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

        else:
            loss.backward()

        if isinstance(x, torch.Tensor):
            grads = x_grad.grad
        else:
            grads = x_grad.grad.cpu().numpy().copy()  # type: ignore

        if not self.all_framework_preprocessing:
            grads = self._apply_preprocessing_gradient(x, grads)

        assert grads.shape == x.shape

        return grads

        Compute the gradient of the loss function w.r.t. `x`.

        :param loss_fn: Loss function w.r.t to which gradient needs to be calculated.
        :param x: Sample input with shape as expected by the model(base image).
        :param y: Sample input with shape as expected by the model(target image).
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.`
        :param layer_name: Name of the layer from which activation needs to be extracted/activation layer.
        :return: Array of gradients of the same shape as `x`.
        
art/estimators/classification/pytorch.py