text: import os
  import urllib
  from matplotlib.pyplot import sca
  import torch
  from torch.utils import model_zoo
  # torch.autograd.set_detect_anomaly(True)
  
  class CheckpointIO(object):
      ''' CheckpointIO class.
  
      It handles saving and loading checkpoints.
  
      Args:
          checkpoint_dir (str): path where checkpoints are saved
      '''
  
      def __init__(self, checkpoint_dir='./chkpts', **kwargs):
          self.module_dict = kwargs
          self.checkpoint_dir = checkpoint_dir
          if not os.path.exists(checkpoint_dir):
              os.makedirs(checkpoint_dir)
  
      def register_modules(self, **kwargs):
          ''' Registers modules in current module dictionary.
          '''
          self.module_dict.update(kwargs)
  
      def save(self, filename, **kwargs):
          ''' Saves the current module dictionary.
  
          Args:
              filename (str): name of output file
          '''
          if not os.path.isabs(filename):
              filename = os.path.join(self.checkpoint_dir, filename)
  
          outdict = kwargs
          for k, v in self.module_dict.items():
              outdict[k] = v.state_dict()
          torch.save(outdict, filename)
  
      def load(self, filename):
          '''Loads a module dictionary from local file or url.
  
          Args:
              filename (str): name of saved module dictionary
          '''
          if is_url(filename):
              return self.load_url(filename)
          else:
              return self.load_file(filename)
  
      def load_file(self, filepath, no_reload=False, ignore_keys=[], only_use_keys=None):
          '''Loads a module dictionary from file.
  
          Args:
              filepath (str): file path of saved module dictionary
          '''
  
          assert not ((len(ignore_keys) > 0) and only_use_keys is not None), \
              'please specify at most one in [ckpt_ignore_keys, ckpt_only_use_keys]'
  
          if filepath is not None and filepath != "None":
              ckpts = [filepath]
          else:
              ckpts = sorted_ckpts(self.checkpoint_dir)
  
          print("=> Found ckpts:",
                "{}".format(ckpts) if len(ckpts) < 5 else "...,{}".format(ckpts[-5:]))
  
          if len(ckpts) > 0 and not no_reload:
              ckpt_file = ckpts[-1]
              print('=> Loading checkpoint from local file', ckpt_file)
              state_dict = torch.load(ckpt_file)
  
              if len(ignore_keys) > 0:
                  to_load_state_dict = {}
                  for k in state_dict.keys():
                      if k in ignore_keys:
                          print("=> [ckpt] Ignoring keys: {}".format(k))
                      else:
                          to_load_state_dict[k] = state_dict[k]
              elif only_use_keys is not None:
                  if not isinstance(only_use_keys, list):
                      only_use_keys = [only_use_keys]
                  to_load_state_dict = {}
                  print(state_dict.keys())
                  for k in only_use_keys:
                      print("=> [ckpt] Only use keys: {}".format(k))
                      to_load_state_dict[k] = state_dict[k]
              else:
                  to_load_state_dict = state_dict
  
              scalars = self.parse_state_dict(to_load_state_dict, ignore_keys)
              return scalars
          else:
              return {}
  
      def load_url(self, url):
          '''Load a module dictionary from url.
  
          Args:
              url (str): url to saved model
          '''
          print(url)
          print('=> Loading checkpoint from url...')
          state_dict = model_zoo.load_url(url, progress=True)
          scalars = self.parse_state_dict(state_dict)
          return scalars
  
      def parse_state_dict(self, state_dict, ignore_keys):
          '''Parse state_dict of model and return scalars.
  
          Args:
              state_dict (dict): State dict of model
      '''
  
          for k, v in self.module_dict.items():
              if k in state_dict:
                  v.load_state_dict(state_dict[k])
              else:
                  if k not in ignore_keys:
                      print('Warning: Could not find %s in checkpoint!' % k)
          scalars = {k: v for k, v in state_dict.items()
                     if k not in self.module_dict}
          return scalars
  
  
  def is_url(url):
      scheme = urllib.parse.urlparse(url).scheme
      return scheme in ('http', 'https')
  
  
  def sorted_ckpts(checkpoint_dir):
      ckpts = []
      if os.path.exists(checkpoint_dir):
          latest = None
          final = None
          ckpts = []
          for fname in sorted(os.listdir(checkpoint_dir)):
              fpath = os.path.join(checkpoint_dir, fname)
              if ".pt" in fname:
                  ckpts.append(fpath)
                  if 'latest' in fname:
                      latest = fpath
                  elif 'final' in fname:
                      final = fpath
          if latest is not None:
              ckpts.remove(latest)
              ckpts.append(latest)
          if final is not None:
              ckpts.remove(final)
              ckpts.append(final)
      return ckpts