def fit_generator(
        self,
        generator: DataGenerator,
        nb_epochs: int = 20,
        scheduler: "torch.optim.lr_scheduler._LRScheduler" | None = None,
        **kwargs,
    ):
        """
        Train a model adversarially with TRADES protocol using a data generator.
        See class documentation for more information on the exact procedure.

        :param generator: Data generator.
        :param nb_epochs: Number of epochs to use for trainings.
        :param scheduler: Learning rate scheduler to run at the end of every epoch.
        :param kwargs: Dictionary of framework-specific arguments. These will be passed as such to the `fit` function of
                                  the target classifier.
        """
        import torch

        logger.info("Performing adversarial training with TRADES protocol")

        if (scheduler is not None) and (
            not isinstance(scheduler, torch.optim.lr_scheduler._LRScheduler)
        ):  # pylint: enable=protected-access
            raise ValueError("Invalid Pytorch scheduler is provided for adversarial training.")

        size = generator.size
        batch_size = generator.batch_size
        if size is not None:
            nb_batches = int(np.ceil(size / batch_size))
        else:
            raise ValueError("Size is None.")

        logger.info("Adversarial Training TRADES")

        for i_epoch in trange(nb_epochs, desc="Adversarial Training TRADES - Epochs"):
            start_time = time.time()
            train_loss = 0.0
            train_acc = 0.0
            train_n = 0.0

            for batch_id in range(nb_batches):  # pylint: disable=unused-variable
                # Create batch data
                x_batch, y_batch = generator.get_batch()
                x_batch = x_batch.copy()

                _train_loss, _train_acc, _train_n = self._batch_process(x_batch, y_batch)

                train_loss += _train_loss
                train_acc += _train_acc
                train_n += _train_n

            if scheduler:
                scheduler.step()

            train_time = time.time()
            logger.info(
                "epoch: %s time(s): %.1f loss: %.4f acc: %.4f",
                i_epoch,
                train_time - start_time,
                train_loss / train_n,
                train_acc / train_n,
            )

        Train a model adversarially with TRADES protocol using a data generator.
        See class documentation for more information on the exact procedure.

        :param generator: Data generator.
        :param nb_epochs: Number of epochs to use for trainings.
        :param scheduler: Learning rate scheduler to run at the end of every epoch.
        :param kwargs: Dictionary of framework-specific arguments. These will be passed as such to the `fit` function of
                                  the target classifier.
        
art/defences/trainer/adversarial_trainer_trades_pytorch.py