code: def wrap_tqdm(
      gen,
      mode,
      n_epochs,
      val_size,
      dataset_size,
      current_epoch_offset=0,
      total_epoch_offset=0,
      skip_val=False,
      quiet=False,
      update_hook=None,
  ):
      assert mode in {"train", "predict", "evaluate"}
      if mode == "predict":
          return gen  # tqdm is handled elsewhere (not sure why)
  
      try:
          total = len(gen)
      except:
          if mode == "train":
              total = dataset_size
          else:
              total = val_size
  
      epoch = 1
  
      def internal_gen():
          nonlocal epoch
          current_epoch = (epoch - 1) % n_epochs + 1
          it = iter(gen())
          if mode == "train":
              desc = "Epoch {}/{}".format(
                  current_epoch + current_epoch_offset,
                  n_epochs + total_epoch_offset,
              )
          else:
              desc = "Validation"
          if skip_val:
              for _, i in zip(range(val_size), it):
                  yield i
          for i in ProgressBar(
              it,
              desc=desc,
              total=total,
              miniters=1,
              leave=current_epoch == n_epochs and mode == "train",
              update_hook=update_hook,
              quiet=quiet,
              current_epoch=current_epoch + current_epoch_offset,
              total_epochs=n_epochs + total_epoch_offset,
          ):
              yield i
          if mode == "train":
              epoch += 1
  
      return internal_gen
  
file_path: finetune/finetune/util/input_utils.py
imports: import math
  import tensorflow
  from finetune.util.timing import ProgressBar
label_desc: Use image node_id 6 for calling a global function with example usage: wrap_tqdm(gen, mode, n_epochs, val_size, dataset_size, current_epoch_offset, total_epoch_offset, skip_val, quiet, update_hook) and returns: internal_gen, gen