class_docstr: Create modules/tensors w. specific devices and dtypes. Examples:
  
  Create MyModule which consists of many different sub-modules and parameters. In this case we can create
  MyModule as a collection of 'meta' tensors by passing `device='meta'` or we can create the module _directly_
  on a CUDA device by passing `device=f'cuda:{local_rank}'` (where `local_rank` is the local GPU id.
  
  with OnDevice(dtype=torch.float16, device='meta'):
      model = MyModel()
  
  with OnDevice(dtype=torch.float16, device=f'cuda:{local_rank}'):
      model = MyModel()
class_docstr_tok: ["Create","modules\/tensors","w.","specific","devices","and","dtypes",".","Examples",":","Create","MyModule","which","consists","of","many","different","sub-modules","and","parameters",".","In","this","case","we","can","create","MyModule","as","a","collection","of","'meta","'","tensors","by","passing","`","device='meta","'","`","or","we","can","create","the","module","_directly_","on","a","CUDA","device","by","passing","`","device=f'cuda",":","{","local_rank","}","'","`","(","where","`","local_rank","`","is","the","local","GPU","id",".","with","OnDevice","(","dtype=torch.float16",",","device='meta","'",")",":","model","=","MyModel","(",")","with","OnDevice","(","dtype=torch.float16",",","device=f'cuda",":","{","local_rank","}","'",")",":","model","=","MyModel","(",")"]
code: def __init__(self, dtype, device="meta", enabled=True):
      self.dtype = dtype
      self.enabled = enabled
      self.device = device
  
      if device == "meta":
          if pkg_version.parse("1.10") > pkg_version.parse(
              torch.__version__
          ):
              raise NotImplementedError(
                  "Meta tensor support is not available, please upgrade to torch 1.10+"
              )
  
file_path: DeepSpeed/deepspeed/utils/init_on_device.py
imports: import torch
  from typing import Callable
  from torch import Tensor
  from packaging import version
label_desc: Use image node_id 1 to create a new OnDevice object from inherited base classes: object with example: obj = OnDevice(dtype, device, enabled)