parameters: (self, remote_device_str: str)
return_statement: return enable_moving_cpu_tensors_to_cuda
docstring: Prepares the initializaiton and returns whether to enable automatically moving CPU tensors to CUDA devices.
docstring_summary: Prepares the initializaiton and returns whether to enable automatically moving CPU tensors to CUDA devices.
docstring_tokens: Prepares the initializaiton and returns whether to enable automatically moving CPU tensors to CUDA devices .
function: def _prepare_init(self, remote_device_str: str) -> bool:
          """
          Prepares the initializaiton and returns whether to enable automatically moving CPU tensors to CUDA devices.
          """
          # Sanity check.
          assert rpc._is_current_rpc_agent_set(), "RemoteModule only works in RPC."
  
          remote_device = _remote_device(remote_device_str)
          self.on = remote_device.worker_name() if remote_device.worker_name() is not None else remote_device.rank()
          self.device = str(remote_device.device())
          agent = rpc._get_current_rpc_agent()
          # If the device map of the remote worker is set,
          # then enable moving any input CPU tensors to the same cuda device.
          self.is_device_map_set = bool(
              agent._get_device_map(agent.get_worker_info(self.on))  # type: ignore[arg-type]
          )
          # ``enable_moving_cpu_tensors_to_cuda`` is less strict than ``is_device_map_set``:
          # If ``enable_moving_cpu_tensors_to_cuda`` is true, but the device map is not set,
          # then any CPU tensors can still be moved to a cuda device to run forward,
          # but the output must be moved back to CPU before being sent over the wire.
          enable_moving_cpu_tensors_to_cuda = torch.device(self.device).type == "cuda"
          return enable_moving_cpu_tensors_to_cuda
function_tokens: def _prepare_init ( self , remote_device_str : str ) -> bool : # Sanity check. assert rpc . _is_current_rpc_agent_set ( ) , "RemoteModule only works in RPC." remote_device = _remote_device ( remote_device_str ) self . on = remote_device . worker_name ( ) if remote_device . worker_name ( ) is not None else remote_device . rank ( ) self . device = str ( remote_device . device ( ) ) agent = rpc . _get_current_rpc_agent ( ) # If the device map of the remote worker is set, # then enable moving any input CPU tensors to the same cuda device. self . is_device_map_set = bool ( agent . _get_device_map ( agent . get_worker_info ( self . on ) ) # type: ignore[arg-type] ) # ``enable_moving_cpu_tensors_to_cuda`` is less strict than ``is_device_map_set``: # If ``enable_moving_cpu_tensors_to_cuda`` is true, but the device map is not set, # then any CPU tensors can still be moved to a cuda device to run forward, # but the output must be moved back to CPU before being sent over the wire. enable_moving_cpu_tensors_to_cuda = torch . device ( self . device ) . type == "cuda" return enable_moving_cpu_tensors_to_cuda