def group_by_bucket(self) -> dict:
        """
        Group the dataset samples into buckets.
        This method will set `self._cached_bucket_sample_dict` to the bucket sample dict.

        Returns:
            dict: a dictionary with bucket id as key and a list of sample indices as value
        """
        if self._cached_bucket_sample_dict is not None:
            return self._cached_bucket_sample_dict, self._cached_num_total_batch

        # use pandarallel to accelerate bucket processing
        log_message("Building buckets using %d workers...", self.num_bucket_build_workers)
        bucket_ids = None
        if dist.get_rank() == 0:
            data = self.dataset.data.copy(deep=True)
            data["id"] = data.index
            bucket_ids = data.parallel_apply(
                apply,
                axis=1,
                method=self.bucket.get_bucket_id,
                seed=self.seed + self.epoch,
                num_bucket=self.bucket.num_bucket,
                fps_max=self.dataset.fps_max,
            )
        dist.barrier()
        bucket_ids = sync_object_across_devices(bucket_ids)
        dist.barrier()

        # group by bucket
        # each data sample is put into a bucket with a similar image/video size
        bucket_sample_dict = defaultdict(list)
        bucket_ids_np = np.array(bucket_ids)
        valid_indices = np.where(bucket_ids_np != None)[0]
        for i in valid_indices:
            bucket_sample_dict[bucket_ids_np[i]].append(i)

        # cache the bucket sample dict
        self._cached_bucket_sample_dict = bucket_sample_dict

        # num total batch
        num_total_batch = self.print_bucket_info(bucket_sample_dict)
        self._cached_num_total_batch = num_total_batch

        return bucket_sample_dict, num_total_batch

        Group the dataset samples into buckets.
        This method will set `self._cached_bucket_sample_dict` to the bucket sample dict.

        Returns:
            dict: a dictionary with bucket id as key and a list of sample indices as value
        
opensora/datasets/sampler.py