class_docstr: CheckpointIO is the base class for all checkpoint IO classes. It defines the interface for checkpoint IO.
  
  
  Examples:
      >>> from colossalai.checkpoint_io import GeneralCheckpointIO
      >>> checkpoint_io = CheckpointIO()
      >>>
      >>> # load model from checkpoint
      >>> model = checkpoint_io.load_model(model, 'model.pt')
      >>>
      >>> # save model to checkpoint, any distributed tensor is gathered by default
      >>> checkpoint_io.save_model(model, 'model.pt')
      >>>
      >>> # if the model contains distributed tensor, and you don't want to gather it
      >>> # each rank will save its own shard of the distributed tensor
      >>> checkpoint_io.save_model(model, 'model.pt', gather_dtensor=False)
      >>>
      >>> # save model to sharded checkpoints
      >>> checkpoint_io.save_model(model, './checkpoints/', shard=True)
      >>>
      >>> # save model to sharded  and assume we don't want to gather distributed tensors
      >>> checkpoint_io.save_model(model, './checkpoints/', shard=True, gather_dtensor=False)
      >>>
      >>> # Note:
      >>> # 1. we don't support loading from distributed tensors, conversion from distributed tensors
      >>> # checkpoints to full tensor checkpoint should be done offline via our CLI
      >>> # 2. you don't have to specify whether the model is sharded or not when loading the model
      >>> # as it will be automatically detected
      >>>
      >>> # load model from sharded checkpoints
      >>> model = checkpoint_io.load_model(model, './checkpoints/')
      >>>
      >>> # load model from unsharded checkpoints
      >>> model = checkpoint_io.load_model(model, './checkpoints/')
      >>>
      >>> # load optimizer from checkpoint
      >>> optimizer = checkpoint_io.load_optimizer(optimizer, 'optimizer.pt')
      >>>
      >>> # save optimizer to checkpoint
      >>> checkpoint_io.save_optimizer(optimizer, 'optimizer.pt')
docstr: Load optimizer from checkpoint.
  
  Args:
      optimizer (Optimizer): optimizer to be loaded.
      checkpoint (str): checkpoint path. This value is made compatibility with the model checkpoints in the
      prefix (str, optional): A prefix added to parameter and buffer
          names to compose the keys in state_dict. Defaults to None.
      size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.
docstr_tok: ["Load","optimizer","from","checkpoint",".","Args",":","optimizer","(","Optimizer",")",":","optimizer","to","be","loaded",".","checkpoint","(","str",")",":","checkpoint","path",".","This","value","is","made","compatibility","with","the","model","checkpoints","in","the","prefix","(","str",",","optional",")",":","A","prefix","added","to","parameter","and","buffer","names","to","compose","the","keys","in","state_dict",".","Defaults","to","None",".","size_per_shard","(","int",",","optional",")",":","Maximum","size","of","checkpoint","shard","file","in","MB",".","This","is","useful","only","when","``","shard=True","``",".","Defaults","to","1024","."]
code: def load_optimizer(
      self,
      optimizer: Optimizer,
      checkpoint: str,
      prefix: str = None,
      size_per_shard: int = 1024,
  ):
      """
      Load optimizer from checkpoint.
  
      Args:
          optimizer (Optimizer): optimizer to be loaded.
          checkpoint (str): checkpoint path. This value is made compatibility with the model checkpoints in the
          prefix (str, optional): A prefix added to parameter and buffer
              names to compose the keys in state_dict. Defaults to None.
          size_per_shard (int, optional): Maximum size of checkpoint shard file in MB. This is useful only when ``shard=True``. Defaults to 1024.
      """
  
      index_file_exists, index_file_path = has_index_file(checkpoint)
  
      if Path(checkpoint).is_dir() and not index_file_exists:
          # if the checkpoint is a directory and there is no index file, raise error
          raise ValueError(f"Cannot find index file in {checkpoint}")
  
      if index_file_exists:
          # the existence of index file means it is a sharded checkpoint
          self.load_sharded_optimizer(
              optimizer, index_file_path, prefix
          )
      else:
          self.load_unsharded_optimizer(optimizer, checkpoint)
  
imports: from abc import ABC, abstractmethod
  from pathlib import Path
  from typing import Optional, Union
  import torch
  import torch.nn
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from colossalai.interface import ModelWrapper
  from .utils import has_index_file
label_desc: Use image node_id 3 for calling the CheckpointIO obj's underlying member method code with example usage: obj.load_optimizer(optimizer, checkpoint, prefix, size_per_shard) without return types