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 fp_tensor_constructor(
      self, fn: Callable, target_fp_dtype: torch.dtype
  ) -> Callable:
      def wrapped_fn(*args, **kwargs) -> Tensor:
          if kwargs.get("device", None) is None:
              kwargs["device"] = self.device
          tensor: Tensor = fn(*args, **kwargs)
          if tensor.is_floating_point():
              tensor = tensor.to(target_fp_dtype)
          return tensor
  
      return wrapped_fn
  
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 2 for calling the OnDevice obj's underlying member method code with example usage: obj.fp_tensor_constructor(fn, target_fp_dtype) and returns: wrapped_fn, tensor