def fit(
        self,
        x: np.ndarray,
        y: np.ndarray,
        batch_size: int = 128,
        nb_epochs: int = 10,
        training_mode: bool = True,
        drop_last: bool = False,
        scheduler: "torch.optim.lr_scheduler._LRScheduler" | None = None,
        verbose: bool = False,
        **kwargs,
    ) -> None:
        """
        Fit the classifier on the training set `(x, y)`.

        :param x: Training data.
        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                  shape (nb_samples,).
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.
        :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by
                          the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                          the last batch will be smaller. (default: ``False``)
        :param scheduler: Learning rate scheduler to run at the start of every epoch.
        :param verbose: Display the training progress bar.
        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
               and providing it takes no effect.
        """
        import torch
        import torch.nn.functional as F
        from torch.utils.data import TensorDataset, DataLoader

        # Set model mode
        self._model.train(mode=training_mode)

        if self._optimizer is None:  # pragma: no cover
            raise ValueError("An optimizer is needed to train the model, but none for provided")

        y = check_and_transform_label_format(y, nb_classes=self.nb_classes)

        # Apply preprocessing
        x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=True)

        # Check label shape
        y_preprocessed = self.reduce_labels(y_preprocessed)

        # Create dataloader
        x_tensor = torch.from_numpy(x_preprocessed)
        y_tensor = torch.from_numpy(y_preprocessed)
        dataset = TensorDataset(x_tensor, y_tensor)
        dataloader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)

        # Start training
        for epoch in trange(nb_epochs, disable=not verbose):
            warmup_v = min(1.0, (epoch + 1) / self.warmup)

            for x_batch, y_batch in dataloader:
                mini_batch_size = len(x_batch) // self.num_noise_vec

                for mini_batch in range(self.num_noise_vec):
                    # Create mini batch
                    inputs = x_batch[mini_batch * mini_batch_size : (mini_batch + 1) * mini_batch_size]
                    labels = y_batch[mini_batch * mini_batch_size : (mini_batch + 1) * mini_batch_size]

                    # Move inputs to GPU
                    inputs = inputs.to(self.device)
                    labels = labels.to(self.device)

                    noises = [torch.randn_like(inputs) * self.scale for _ in range(self.num_noise_vec)]

                    # Attack and find adversarial examples
                    self._model.eval()
                    inputs, inputs_adv = self._smooth_mix_pgd_attack(inputs, labels, noises, warmup_v)
                    self._model.train(mode=training_mode)

                    in_clean_c = torch.cat([inputs + noise for noise in noises], dim=0)
                    logits_c = self._model(in_clean_c)[-1]
                    labels_c = labels.repeat(self.num_noise_vec)

                    logits_c_chunk = torch.chunk(logits_c, self.num_noise_vec, dim=0)
                    clean_sm = F.softmax(torch.stack(logits_c_chunk), dim=-1)
                    clean_avg_sm = torch.mean(clean_sm, dim=0)
                    loss_xent = F.cross_entropy(logits_c, labels_c, reduction="none")

                    # mix adversarial examples
                    in_mix, labels_mix = self._mix_data(inputs, inputs_adv, clean_avg_sm)
                    in_mix_c = torch.cat([in_mix + noise for noise in noises], dim=0)
                    labels_mix_c = labels_mix.repeat(self.num_noise_vec, 1)
                    logits_mix_c = F.log_softmax(self._model(in_mix_c)[-1], dim=1)

                    preds = torch.argmax(clean_avg_sm, dim=-1)
                    ind_correct = (preds == labels).float()
                    ind_correct = ind_correct.repeat(self.num_noise_vec)

                    loss_mixup = F.kl_div(logits_mix_c, labels_mix_c, reduction="none").sum(1)
                    loss = loss_xent.mean() + self.eta * warmup_v * (ind_correct * loss_mixup).mean()

                    # compute gradient and do SGD step
                    self._optimizer.zero_grad()
                    loss.backward()
                    self._optimizer.step()

            if scheduler is not None:
                scheduler.step()

        Fit the classifier on the training set `(x, y)`.

        :param x: Training data.
        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or index labels of
                  shape (nb_samples,).
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.
        :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by
                          the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                          the last batch will be smaller. (default: ``False``)
        :param scheduler: Learning rate scheduler to run at the start of every epoch.
        :param verbose: Display the training progress bar.
        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
               and providing it takes no effect.
        
art/estimators/certification/randomized_smoothing/smooth_mix/pytorch.py