def update_batchnorm(self, x: np.ndarray, batch_size: int, nb_epochs: int = 1) -> None:
        """
        Method to update the batchnorm of a neural network on small datasets when it was pre-trained

        :param x: Training data.
        :param batch_size: Size of batches.
        :param nb_epochs: How many times to forward pass over the input data
        """
        import torch

        if self.mode != "ViT":
            raise ValueError("Accessing a ViT specific functionality while running in CNN mode")

        self.model.train()

        ind = np.arange(len(x))
        num_batch = int(len(x) / float(batch_size))

        with torch.no_grad():
            for _ in tqdm(range(nb_epochs)):
                for m in tqdm(range(num_batch)):
                    i_batch = self.ablator.forward(
                        np.copy(x[ind[m * batch_size : (m + 1) * batch_size]]), column_pos=random.randint(0, x.shape[3])
                    )
                    _ = self.model(i_batch)

        Method to update the batchnorm of a neural network on small datasets when it was pre-trained

        :param x: Training data.
        :param batch_size: Size of batches.
        :param nb_epochs: How many times to forward pass over the input data
        
art/estimators/certification/derandomized_smoothing/pytorch.py