code: def iter_torch_batches(
      ds: Dataset,
      batch_size: Optional[int] = None,
      local_shuffle_buffer_size: Optional[int] = None,
      prefetch_batches: int = 0,
      use_default_params: bool = False,
  ) -> Dataset:
      num_batches = 0
      if use_default_params:
          for batch in ds.iter_torch_batches():
              num_batches += 1
      else:
          for batch in ds.iter_torch_batches(
              batch_size=batch_size,
              local_shuffle_buffer_size=local_shuffle_buffer_size,
              prefetch_batches=prefetch_batches,
          ):
              num_batches += 1
      print(
          "iter_torch_batches done, block_format:",
          "pyarrow",
          "num_rows:",
          ds.count(),
          "num_blocks:",
          ds.num_blocks(),
          "num_batches:",
          num_batches,
      )
      return ds
  
file_path: ray/release/nightly_tests/dataset/iter_tensor_batches_benchmark.py
imports: import argparse
  import numpy
  from typing import Optional, Union, List
  import ray
  from ray.data.dataset import Dataset
  from benchmark import Benchmark
label_desc: Use image node_id 1 for calling a global function with example usage: iter_torch_batches(ds, batch_size, local_shuffle_buffer_size, prefetch_batches, use_default_params) and returns: ds