def palettize(self, X, cX, pad, partitions) -> _torch.Tensor:
        """
        This method is run during inference time by the forward method of the ``FakePalettize`` class. It calculates the
        weight from the ``LUT`` and ``indices`` across all partitions and returns them.
        """
        batch_partitions = []
        seq_partitions = []
        most_common_numel = self.partition_numel[partitions].mode()[0]

        for p in partitions:
            if self.partition_numel[p] == most_common_numel and self.labels[p] is None:
                batch_partitions.append(p)
            else:
                seq_partitions.append(p)

        if len(batch_partitions) == 1 or not self.palett_batch_mode:
            seq_partitions += batch_partitions
            batch_partitions = []

        if seq_partitions:
            X = self.palettize_seq(X, cX, pad, seq_partitions)

        if batch_partitions:
            X = self.palettize_batch(X, cX, pad, batch_partitions)

        return X

        This method is run during inference time by the forward method of the ``FakePalettize`` class. It calculates the
        weight from the ``LUT`` and ``indices`` across all partitions and returns them.
        
coremltools/optimize/torch/palettization/fake_palettize.py