class_docstr: Booster is a high-level API for training neural networks. It provides a unified interface for
  training with different precision, accelerator, and plugin.
  
  
  ```python
  # Following is pseudocode
  
  colossalai.launch(...)
  plugin = GeminiPlugin(...)
  booster = Booster(precision='fp16', plugin=plugin)
  
  model = GPT2()
  optimizer = HybridAdam(model.parameters())
  dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8)
  lr_scheduler = LinearWarmupScheduler()
  criterion = GPTLMLoss()
  
  model, optimizer, criterion, dataloader, lr_scheduler = booster.boost(model, optimizer, criterion, dataloader, lr_scheduler)
  
  for epoch in range(max_epochs):
      for input_ids, attention_mask in dataloader:
          outputs = model(input_ids.cuda(), attention_mask.cuda())
          loss = criterion(outputs.logits, input_ids)
          booster.backward(loss, optimizer)
          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()
  ```
  
  Args:
      device (str or torch.device): The device to run the training. Default: None.
                                    If plugin is not used or plugin doesn't control the device,
                                    this argument will be set as training device ('cuda' will be used if argument is None).
      mixed_precision (str or MixedPrecision): The mixed precision to run the training. Default: None.
                              If the argument is a string, it can be 'fp16', 'fp16_apex', 'bf16', or 'fp8'.
                              'fp16' would use PyTorch AMP while `fp16_apex` would use Nvidia Apex.
      plugin (Plugin): The plugin to run the training. Default: None.
class_docstr_tok: ["Booster","is","a","high-level","API","for","training","neural","networks",".","It","provides","a","unified","interface","for","training","with","different","precision",",","accelerator",",","and","plugin",".","``","`","python","#","Following","is","pseudocode","colossalai.launch","(","...",")","plugin","=","GeminiPlugin","(","...",")","booster","=","Booster","(","precision='fp16","'",",","plugin=plugin",")","model","=","GPT2","(",")","optimizer","=","HybridAdam","(","model.parameters","(",")",")","dataloader","=","plugin.prepare_dataloader","(","train_dataset",",","batch_size=8",")","lr_scheduler","=","LinearWarmupScheduler","(",")","criterion","=","GPTLMLoss","(",")","model",",","optimizer",",","criterion",",","dataloader",",","lr_scheduler","=","booster.boost","(","model",",","optimizer",",","criterion",",","dataloader",",","lr_scheduler",")","for","epoch","in","range","(","max_epochs",")",":","for","input_ids",",","attention_mask","in","dataloader",":","outputs","=","model","(","input_ids.cuda","(",")",",","attention_mask.cuda","(",")",")","loss","=","criterion","(","outputs.logits",",","input_ids",")","booster.backward","(","loss",",","optimizer",")","optimizer.step","(",")","lr_scheduler.step","(",")","optimizer.zero_grad","(",")","``","`","Args",":","device","(","str","or","torch.device",")",":","The","device","to","run","the","training",".","Default",":","None",".","If","plugin","is","not","used","or","plugin","does","n't","control","the","device",",","this","argument","will","be","set","as","training","device","(","'cuda","'","will","be","used","if","argument","is","None",")",".","mixed_precision","(","str","or","MixedPrecision",")",":","The","mixed","precision","to","run","the","training",".","Default",":","None",".","If","the","argument","is","a","string",",","it","can","be","'fp16","'",",","'fp16_apex","'",",","'bf16","'",",","or","'fp8","'",".","'fp16","'","would","use","PyTorch","AMP","while","`","fp16_apex","`","would","use","Nvidia","Apex",".","plugin","(","Plugin",")",":","The","plugin","to","run","the","training",".","Default",":","None","."]
docstr: Load model from checkpoint.
  
  Args:
      model (nn.Module or ModelWrapper): A model boosted by Booster.
      checkpoint (str): Path to the checkpoint. It must be a local path.
          It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.
      strict (bool, optional): whether to strictly enforce that the keys
          in :attr:`state_dict` match the keys returned by this module's
          :meth:`~torch.nn.Module.state_dict` function. Defaults to True.
docstr_tok: ["Load","model","from","checkpoint",".","Args",":","model","(","nn.Module","or","ModelWrapper",")",":","A","model","boosted","by","Booster",".","checkpoint","(","str",")",":","Path","to","the","checkpoint",".","It","must","be","a","local","path",".","It","should","be","a","directory","path","if","the","checkpoint","is","sharded",".","Otherwise",",","it","should","be","a","file","path",".","strict","(","bool",",","optional",")",":","whether","to","strictly","enforce","that","the","keys","in",":","attr",":","`","state_dict","`","match","the","keys","returned","by","this","module's",":","meth",":","`","~torch.nn.Module.state_dict","`","function",".","Defaults","to","True","."]
code: def load_model(
      self,
      model: Union[nn.Module, ModelWrapper],
      checkpoint: str,
      strict: bool = True,
  ) -> None:
      """Load model from checkpoint.
  
      Args:
          model (nn.Module or ModelWrapper): A model boosted by Booster.
          checkpoint (str): Path to the checkpoint. It must be a local path.
              It should be a directory path if the checkpoint is sharded. Otherwise, it should be a file path.
          strict (bool, optional): whether to strictly enforce that the keys
              in :attr:`state_dict` match the keys returned by this module's
              :meth:`~torch.nn.Module.state_dict` function. Defaults to True.
      """
      self.checkpoint_io.load_model(model, checkpoint, strict)
  
imports: import warnings
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, Iterator, List, Optional, Union
  import torch
  import torch.nn
  from torch.optim import Optimizer
  from torch.optim.lr_scheduler import _LRScheduler
  from torch.utils.data import DataLoader
  import colossalai.interface.pretrained
  from colossalai.checkpoint_io import GeneralCheckpointIO
  from colossalai.interface import ModelWrapper, OptimizerWrapper
  from .accelerator import Accelerator
  from .mixed_precision import MixedPrecision, mixed_precision_factory
  from .plugin import Plugin
  from .plugin.pp_plugin_base import PipelinePluginBase
label_desc: Use image node_id 6 for calling the Booster obj's underlying member method code with example usage: obj.load_model(model, checkpoint, strict) without return types