def forward(
        self, x: "torch.Tensor", y: "torch.Tensor" | None = None
    ) -> tuple["torch.Tensor", "torch.Tensor" | None]:
        """
        Apply CutMix data augmentation to sample `x`.

        :param x: Sample to augment with shape of `NCHW`, `NHWC`, `NCFHW` or `NFHWC`.
        :param y: Labels of `x` either one-hot or multi-hot encoded of shape `(nb_samples, nb_classes)`
                  or class indices of shape `(nb_samples,)`.
        :return: Data augmented sample. The returned labels will be probability vectors of shape
                 `(nb_samples, nb_classes)`.
        """
        import torch

        if y is None:
            raise ValueError("Labels `y` cannot be None.")

        # convert labels to one-hot encoding
        if len(y.shape) == 2:
            y_one_hot = y
        elif len(y.shape) == 1:
            y_one_hot = torch.nn.functional.one_hot(y, self.num_classes)
        else:
            raise ValueError(
                "Shape of labels not recognised. "
                "Please provide labels in shape (nb_samples,) or (nb_samples, nb_classes)"
            )

        x_ndim = len(x.shape)

        # NCHW/NHWC/NFHWC --> NCFHW
        if x_ndim == 4:
            if self.channels_first:
                # NCHW --> NCFHW
                x_ncfhw = torch.unsqueeze(x, dim=2)
            else:
                # NHWC --> NCHW --> NCFHW
                x_nchw = torch.permute(x, (0, 3, 1, 2))
                x_ncfhw = torch.unsqueeze(x_nchw, dim=2)
        elif x_ndim == 5:
            if self.channels_first:
                # NCFHW
                x_ncfhw = x
            else:
                # NFHWC --> NCFHW
                x_ncfhw = torch.permute(x, (0, 4, 1, 2, 3))
        else:
            raise ValueError("Unrecognized input dimension. CutMix can only be applied to image and video data.")

        n, _, _, height, width = x_ncfhw.shape
        x_aug = x_ncfhw.clone()
        y_aug = y_one_hot.clone().float()

        # sample the combination ratio from the Beta distribution
        lmb = np.random.beta(self.alpha, self.alpha)
        cut_ratio = np.sqrt(1.0 - lmb)
        cut_height = int(height * cut_ratio)
        cut_width = int(width * cut_ratio)

        # randomly choose indices for samples to mix
        indices = torch.randperm(n, device=x.device)

        # generate a random bounding box per image
        for idx1, idx2 in enumerate(tqdm(indices, desc="CutMix", disable=not self.verbose)):
            prob = np.random.rand()
            if prob < self.probability:
                # uniform sampling
                center_x = torch.randint(0, height, (1,))
                center_y = torch.randint(0, width, (1,))
                bby1 = torch.clamp(center_y - cut_height // 2, 0, height)
                bbx1 = torch.clamp(center_x - cut_width // 2, 0, width)
                bby2 = torch.clamp(center_y + cut_height // 2, 0, height)
                bbx2 = torch.clamp(center_x + cut_width // 2, 0, width)

                # insert image bounding box
                x_aug[idx1, :, :, bbx1:bbx2, bby1:bby2] = x_ncfhw[idx2, :, :, bbx1:bbx2, bby1:bby2]  # type: ignore
                # mix labels
                y_aug[idx1] = lmb * y_aug[idx1] + (1.0 - lmb) * y_one_hot[idx2]  # type: ignore

        x_ncfhw = x_aug

        # NCHW/NHWC/NFHWC <-- NCFHW
        if x_ndim == 4:
            if self.channels_first:
                # NCHW <-- NCFHW
                x_aug = torch.squeeze(x_ncfhw, dim=2)
            else:
                # NHWC <-- NCHW <-- NCFHW
                x_nchw = torch.squeeze(x_ncfhw, dim=2)
                x_aug = torch.permute(x_nchw, (0, 2, 3, 1))
        elif x_ndim == 5:
            if self.channels_first:
                # NCFHW
                x_aug = x_ncfhw
            else:
                # NFHWC <-- NCFHW
                x_aug = torch.permute(x_ncfhw, (0, 2, 3, 4, 1))

        return x_aug, y_aug

        Apply CutMix data augmentation to sample `x`.

        :param x: Sample to augment with shape of `NCHW`, `NHWC`, `NCFHW` or `NFHWC`.
        :param y: Labels of `x` either one-hot or multi-hot encoded of shape `(nb_samples, nb_classes)`
                  or class indices of shape `(nb_samples,)`.
        :return: Data augmented sample. The returned labels will be probability vectors of shape
                 `(nb_samples, nb_classes)`.
        
art/defences/preprocessor/cutmix/cutmix_pytorch.py