code: def lms_scheduler(
      ctx, use_lms, gpu_memory_size=None, window_length=None
  ):
      _check_list = [x.split(":")[0] for x in ctx.backend]
      if "cudnn" not in _check_list and "cuda" not in _check_list:
          logger.warning(
              "ctx passed to scheduler doesn't have cuda/cudnn backend. lms scheduler will not be used."
          )
          use_lms = False
  
      comm = current_communicator()
      if comm:
          logger.log(
              99,
              f"[OoC] Currently OoC is disabled for Multi-GPU training.",
          )
          use_lms = False
  
      if use_lms:
          gpu_index = 0
          if "cuda" in str(ctx.backend):
              gpu_index = int(ctx.device_id)
          else:
              logger.log(
                  99, f"[OoC] OoC is only enabled for GPU training."
              )
              raise Exception
  
          if gpu_memory_size is None or gpu_memory_size == 0:
              try:
                  handle = nvml.nvmlDeviceGetHandleByIndex(gpu_index)
                  total_memory = nvml.nvmlDeviceGetMemoryInfo(
                      handle
                  ).total
                  gpu_memory_size = int(total_memory * 0.7)
              except:
                  logger.log(
                      99,
                      f"[OoC] Could not get GPU memory size using default value(6GB).",
                  )
                  gpu_memory_size = 6e9  # default 6 GiB
                  pass
  
          if window_length is None or window_length == 0:
              window_length = int(gpu_memory_size * 1.5)
  
          logger.log(
              99,
              f"[OoC] gpu_memory_limit: {gpu_memory_size / 1e9}GB, prefetch_window_length: {window_length / 1e9}GB",
          )
          # Change array preference so that lms works well.
          # import nnabla_ext.cuda.init as cuda_init
          # cuda_init.prefer_cpu_pinned_array()
          # cuda_init.prefer_cuda_virtual_array()
          from nnabla.ext_utils import get_extension_context
  
          be, tc = ctx.backend[0].split(":")
          cpu_ctx = get_extension_context(
              "cpu", device_id="", type_config=tc
          )
          return SwapInOutScheduler(
              cpu_ctx, ctx, gpu_memory_size, window_length
          )
      else:
  
          class DummyScheduler(object):
              function_pre_hook = None
              function_post_hook = None
              update_pre_hook = None
              update_post_hook = None
  
              def start_scheduling(self):
                  return None
  
              def end_scheduling(self):
                  return None
  
              def __enter__(self):
                  return self
  
              def __exit__(self, exc_type, exc_val, exc_tb):
                  pass
  
          return DummyScheduler()
  
imports: import os
  import re
  import timeit
  from contextlib import contextmanager
  import nnabla
  import nnabla.utils.callback
  import numpy
  from nnabla.config import nnabla_config
  from nnabla.logger import logger
  from nnabla.utils import nnabla_pb2
  from nnabla.utils.communicator_util import current_communicator
  from nnabla.utils.get_file_handle import get_file_handle_save, get_decorated_file_loader, load_files
  from contextlib import contextmanager
  import nnabla.logger
  from nnabla.lms import SwapInOutScheduler
label_desc: Use image node_id 12 for calling a global function with example usage: lms_scheduler(ctx, use_lms, gpu_memory_size, window_length) and returns: SwapInOutScheduler, DummyScheduler, None, None, self