class_docstr: Model used to check cpu offload memory saving
class_docstr_tok: ["Model","used","to","check","cpu","offload","memory","saving"]
code: def __init__(self, enable_checkpoint=False, cpu_offload=False):
      super().__init__()
  
      torch.manual_seed(0)  # make sure weights are deterministic.
  
      # These numbers are picked to show cpu_offload memory saving.
      # Inner (recomputed) activation sizes need to be just right
      # to show the benefit.
      self.layers = nn.Sequential(
          nn.Sequential(
              nn.Linear(4, 4), nn.Linear(4, 4), nn.Linear(4, 8)
          ),
          nn.Sequential(
              nn.Linear(8, 4), nn.Linear(4, 4), nn.Linear(4, 4)
          ),
          nn.Sequential(
              nn.Linear(4, 6), nn.Linear(6, 8), nn.Linear(8, 2)
          ),
      )
  
      if enable_checkpoint:
          for i, layer in enumerate(self.layers):
              # Only middle layer needs to have offloading
              self.layers[i] = checkpoint_wrapper(
                  layer, cpu_offload if i == 1 else False
              )
  
imports: import pytest
  import torch
  import torch.nn
  from torch.utils.checkpoint import checkpoint
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper, disable_checkpointing
  from fairscale.nn.misc import FlattenParamsWrapper
  from fairscale.nn.misc import checkpoint_wrapper
label_desc: Use image node_id 1 to create a new CpuOffloadModel object from inherited base classes: nn with example: obj = CpuOffloadModel(enable_checkpoint, cpu_offload)