class_docstr: Wraps an arbitrary :class:`nn.Sequential <torch.nn.Sequential>` module
  to train by offloading majority of the model parameters to the CPU.
  `OffloadModel` is heavily inspired by the _L2L algorithm and _Zero-Offload.
  ::
  
      model = get_model()
      offload_model = OffloadModel(model, device,
                                  offload_device=torch.device(“cpu”),
                                  num_slices=3,
                                  checkpoint_activation=True,
                                  num_microbatches=5)
  
  .. _L2L: https://arxiv.org/abs/2002.05645
  .. _Zero-Offload: https://arxiv.org/abs/2101.06840
  
  At each step, a layer(or series of layers) are loaded
  onto the GPU for the forward and backward pass with intermediate
  activations being copied onto the GPU as required. Once the forward
  or backward pass is completed for a given shard, it is moved back to
  the CPU again.
  
  `OffloadModel` supports activation checkpointing which reduces
  the memory footprint. You can also increase the number of
  microbatches which translates to more computation cycles for
  every shard load. This helps offset the cost of moving the shard
  from the CPU to GPU and vice versa.
  
  Note: OffloadModel currently only supports nn.Sequential models.
  
  Args:
      module (~torch.nn.Sequential): Module to be offloaded.
  
      device (torch.device):
          Device where the active model should reside.
  
      offload_device (torch.device):
          Device where the inactive model should reside.
  
      num_slices (int):
          Number of slices into which the model should be chunked.
  
      checkpoint_activation (bool):
          Boolean to indicate if we want to checkpoint intermediate
          activation states on the CPU. Default value is False.
  
      num_microbatches (int):
          Number of microbatches which should be run per model
          shard on device.
code: def __init__(
      self,
      model: Any,
      device: torch.device,
      offload_device: torch.device = torch.device("cpu"),
      num_slices: int = 3,
      checkpoint_activation: bool = False,
      num_microbatches: int = 1,
  ):
      super().__init__()
      if not model:
          raise TypeError(
              "`model` argument to `OffloadModel` cannot be None."
          )
  
      if not device:
          raise TypeError(
              "`device` argument to `OffloadModel` cannot be None."
          )
  
      if not (isinstance(model, nn.Sequential) or type(model) == list):
          raise TypeError(
              "`model` argument to `OffloadModel` must be of type `nn.Sequential`."
          )
  
      if not torch.cuda.is_available():
          raise TypeError(
              "CUDA must be available as one of the compute devices for `OffloadModel`."
          )
  
      self.device = device
      self.offload_device = offload_device
      # List of model shards that will be placed on/off the device.
      self.model_slices: List[nn.Module] = []
  
      # TODO(anj): Add an experimental flag for using this instead of modifying the
      # arg type.
      if type(model) == list:
          # This is already sharded using the auto shard functinality.
          for i, m in enumerate(model):
              self.model_slices.append(
                  ModelShard(
                      cpu_model_shard=m,
                      device=device,
                      offload_device=offload_device,
                      index=i,
                  )
              )
      else:
          # Slice the model into roughly equivalent sequential shards.
          splits = _split(model, num_slices)  # type: ignore
  
          for i, split in enumerate(splits):
              # Add one model handling this slice
              self.model_slices.append(
                  ModelShard(
                      cpu_model_shard=nn.Sequential(*split),
                      device=device,
                      offload_device=offload_device,
                      index=i,
                  )
              )
  
      # Expose a unified view of the slices
      self._model = torch.nn.Sequential(*self.model_slices)
  
      # intermediate activations at the slice boundaries.
      self._activations: List[Tuple] = []
  
      # Currently we only support microbatches with activation checkpointing.
      if not checkpoint_activation and num_microbatches > 1:
          raise RuntimeError(
              "We currently only support microbatches with activation checkpointing."
          )
  
      # Bool indicating if we want to checkpoint activation on the host.
      self._checkpoint_activation = checkpoint_activation
  
      # Number of microbatches to run per batch on the device
      self._num_microbatches = num_microbatches
  
imports: from builtins import isinstance
  import functools
  import logging
  from typing import Any, List, Tuple
  import torch
  from torch import nn
label_desc: Use image node_id 1 to create a new OffloadModel object from inherited base classes: nn with example: obj = OffloadModel(model, device, offload_device, num_slices, checkpoint_activation, num_microbatches)