class_docstr: Wrap one shard of the model, make it possible to load parameters on the
  fly for the FW and BW pass on the given device.
code: def __init__(
      self,
      cpu_model_shard: nn.Module,
      device: torch.device,
      offload_device: torch.device,
      index: int,
  ):
      super().__init__()
      self.model_shard = cpu_model_shard
      self.index = index
  
      # Save all the parameter sizes to be able to restore them
      self.device = device
      torch.cuda.device(self.device)
  
      self.offload_device = offload_device
  
      self.model_shard.to(offload_device)
      self._cpu_to_gpu_stream = torch.cuda.Stream(device=self.device)
      self._gpu_to_cpu_stream = torch.cuda.Stream(device=self.device)
  
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 ModelShard object from inherited base classes: nn with example: obj = ModelShard(cpu_model_shard, device, offload_device, index)