code: def __init__(
      self,
      data_source,
      batch_size: int,
      drop_last: bool,
      shuffle: bool = True,
  ) -> None:
      if isinstance(next(iter(data_source)), dict):
          first_key = next(iter(next(iter(data_source)).keys()))
          self.lengths = [len(d[first_key]) for d in data_source]
      else:
          self.lengths = [len(d) for d in data_source]
      self.batch_size = batch_size
      self.drop_last = drop_last
      self.shuffle = shuffle
  
file_path: llama-recipes/src/llama_recipes/data/sampler.py
imports: import random
  from itertools import islice
  import numpy
  import torch
label_desc: Use image node_id 1 to create a new LengthBasedBatchSampler object from inherited base classes: torch.utils.data with example: obj = LengthBasedBatchSampler(data_source, batch_size, drop_last, shuffle)