code: def get_iterator_k_split(
      batch: List[torch.Tensor], num_microbatches: int
  ) -> Iterator:
      if isinstance(batch, dict):
          items = list(batch.items())
          assert (
              items[0][1].shape[0] % num_microbatches == 0
          ), "Issue with batch size configuration!"
          split_batch = [
              torch.tensor_split(item[1], num_microbatches, dim=0)
              for item in items
          ]
          microbatches = [
              [
                  (items[i][0], split_batch[i][j])
                  for i in range(len(items))
              ]
              for j in range(num_microbatches)
          ]
          microbatches = [dict(elem) for elem in microbatches]
      else:
          assert (
              batch[0].shape[0] % num_microbatches == 0
          ), "Issue with batch size configuration!"
          split_batch = [
              torch.tensor_split(item, num_microbatches, dim=0)
              if torch.is_tensor(item)
              else item
              for item in batch
          ]
          microbatches = [
              [
                  elem[i] if elem is not None else elem
                  for elem in split_batch
              ]
              for i in range(num_microbatches)
          ]
  
      return itertools.chain(microbatches)
  
imports: import itertools
  import math
  from typing import Dict, Iterator, List, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 25 for calling a global function with example usage: get_iterator_k_split(batch, num_microbatches) and returns: itertools