code: def _cache_images(num_images):
      cache = {"images": None, "done": False}
  
      def decorator(fun):
          def step(state):
              if state[torchbearer.BATCH] == 0:
                  cache["done"] = False
  
              if not cache["done"]:
                  data = fun(state)
  
                  if cache["images"] is None:
                      remaining = (
                          num_images
                          if num_images < data.size(0)
                          else data.size(0)
                      )
  
                      cache["images"] = data[:remaining]
                  else:
                      remaining = num_images - cache["images"].size(0)
  
                      if remaining > data.size(0):
                          remaining = data.size(0)
  
                      cache["images"] = torch.cat(
                          (cache["images"], data[:remaining]), dim=0
                      )
  
                  if cache["images"].size(0) >= num_images:
                      res = cache["images"]
                      cache["done"] = True
                      cache["images"] = None
                      return res
  
          return step
  
      return decorator
  
imports: import torchbearer
  from torchbearer import Callback
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: _cache_images(num_images) and returns: decorator, step, res