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 __exit__(self, exc_type, exc_value, traceback):
      if not self.enabled:
          return
      torch.Tensor.__new__ = torch.Tensor.__old_new__
      torch.empty = self._orig_torch_empty
      torch.zeros = self._orig_torch_zeros
      torch.ones = self._orig_torch_ones
      torch.full = self._orig_torch_full
  
file_path: H2O/h2o_flexgen/benchmark/third_party/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 5 for calling the OnDevice obj's underlying member method code with example usage: obj.__exit__(exc_type, exc_value, traceback) without return types