def fit(
        self,
        x: np.ndarray,
        y: np.ndarray,
        batch_size: int = 128,
        nb_epochs: int = 20,
        train_dict: dict | None = None,
        **kwargs,
    ) -> None:
        """
        Fit the classifier on the training set `(x, y)`. Each classifier will be trained with the
        same parameters unless train_dict is provided. If train_dict is provided, the model id's
        specified will use the training parameters in train_dict instead.

        :param x: Training data.
        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes).
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param train_dict: A dictionary of training args if certain models need specialized arguments.
               The key should be the model's partition id and this will override any default training
               parameters including batch_size and nb_epochs.
        :param kwargs: Dictionary of framework-specific arguments.
        """
        if self.can_fit:
            # First, partition the data using the hash function
            partition_ind = [[] for _ in range(self.ensemble_size)]  # type: list[list[int]]
            for i, p_x in enumerate(x):
                partition_id = int(self.hash_function(p_x))
                partition_ind[partition_id].append(i)

            # Then, train each model on its assigned partition
            for i in range(self.ensemble_size):
                current_x = x[np.array(partition_ind[i])]
                current_y = y[np.array(partition_ind[i])]

                if train_dict is not None and i in train_dict.keys():
                    self.classifiers[i].fit(current_x, current_y, **train_dict[i])
                else:
                    self.classifiers[i].fit(current_x, current_y, batch_size=batch_size, nb_epochs=nb_epochs, **kwargs)
        else:
            warnings.warn("Cannot call fit() for an ensemble of pre-trained classifiers.")

        Fit the classifier on the training set `(x, y)`. Each classifier will be trained with the
        same parameters unless train_dict is provided. If train_dict is provided, the model id's
        specified will use the training parameters in train_dict instead.

        :param x: Training data.
        :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes).
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param train_dict: A dictionary of training args if certain models need specialized arguments.
               The key should be the model's partition id and this will override any default training
               parameters including batch_size and nb_epochs.
        :param kwargs: Dictionary of framework-specific arguments.
        
art/estimators/classification/deep_partition_ensemble.py