code: def batch(data, batch_size, batch_size_fn=None):
      """Yield elements from data in chunks of batch_size."""
      if batch_size_fn is None:
  
          def batch_size_fn(new, count, sofar):
              return count
  
      minibatch, size_so_far = [], 0
      for ex in data:
          minibatch.append(ex)
          size_so_far = batch_size_fn(ex, len(minibatch), size_so_far)
          if size_so_far == batch_size:
              yield minibatch
              minibatch, size_so_far = [], 0
          elif size_so_far > batch_size:
              yield minibatch[:-1]
              minibatch, size_so_far = minibatch[-1:], batch_size_fn(
                  ex, 1, 0
              )
      if minibatch:
          yield minibatch
  
imports: import logging
  import math
  import random
  import torch
  from torchtext.data.utils import RandomShuffler
  from .batch import Batch
  from .dataset import Dataset
label_desc: Use image node_id 1 for calling a global function with example usage: batch(data, batch_size, batch_size_fn) and returns: count