code: def __init__(
      self,
      dataset,
      lengths,
      max_tokens,
      num_buckets,
      shuffle=False,
      batch_size=None,
  ):
      super().__init__()
  
      assert len(dataset) == len(lengths)
  
      self.dataset = dataset
  
      max_length = max(lengths)
      min_length = min(lengths)
  
      assert max_tokens >= max_length
  
      buckets = torch.linspace(min_length, max_length, num_buckets)
      lengths = torch.tensor(lengths)
      bucket_assignments = torch.bucketize(lengths, buckets)
  
      idx_length_buckets = [
          (idx, length, bucket_assignments[idx])
          for idx, length in enumerate(lengths)
      ]
      if shuffle:
          idx_length_buckets = random.sample(
              idx_length_buckets, len(idx_length_buckets)
          )
      else:
          idx_length_buckets = sorted(
              idx_length_buckets, key=lambda x: x[1], reverse=True
          )
  
      sorted_idx_length_buckets = sorted(
          idx_length_buckets, key=lambda x: x[2]
      )
      self.batches = _batch_by_token_count(
          [
              (idx, length)
              for idx, length, _ in sorted_idx_length_buckets
          ],
          max_tokens,
          batch_size=batch_size,
      )
  
imports: import os
  import random
  import torch
  import torchaudio
  from pytorch_lightning import LightningDataModule
label_desc: Use image node_id 1 to create a new CustomBucketDataset object from inherited base classes: torch.utils.data with example: obj = CustomBucketDataset(dataset, lengths, max_tokens, num_buckets, shuffle, batch_size)