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 forward(self, *inputs: Any, **_: Any) -> Any:
      # `apply` calls the `forward` function of the `OffloadFunction` class
      # and the `forward` function calls `inputs` on the first model shard.
      # Please see https://pytorch.org/docs/stable/autograd.html#function for more details.
  
      # We need the second param to be a dummy input to enable the
      # backward pass to be triggered for integer inputs.
      if self._checkpoint_activation:
          return OffloadFunction.apply(
              *inputs, torch.tensor([], requires_grad=True), self
          )
  
      self._activations = []
      for index in range(-1, len(self.model_slices)):
          if index >= 0:
              # TODO(anj-s): This might be a redundant call since we have the previous
              # activation on the device already.
              self._activations[index] = tuple(
                  [a.cuda() for a in list(self._activations[index])]
              )
              inputs = self._activations[index]
              inputs = self.model_slices[index](*inputs)
          # Call the custom autograd hooks (discard/load slices FW and BW)
          inputs = ShardSyncLayer.apply(
              inputs, index, self.model_slices, self
          )
          self._activations.append(inputs)
          if index >= 0:
              self._activations[index] = tuple(
                  [a.cpu() for a in list(self._activations[index])]
              )
  
      result = self._activations[-1]
      result = tuple([r.cuda() for r in result])
      return result[0] if len(result) == 1 else result
  
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 2 for calling the OffloadModel obj's underlying member method code with example usage: obj.forward() and returns: unknown, OffloadFunction