code: def init(
      self,
      module: nn.Module,
      ignored_modules: Set[nn.Module],
      **kwargs,
  ) -> None:
      if _is_fully_sharded(module):
          raise RuntimeError(
              "Cannot apply `replicate()` on a Module already managed by `fully_shard`"
          )
  
      if self.has_initialized:
          return
  
      self.has_initialized = True
      self.module = module
      ignored_params = {
          p for m in ignored_modules for p in m.parameters()
      }
      self._collect_params(module, ignored_modules, ignored_params)
      module.register_forward_pre_hook(
          self.forward_pre_hook, with_kwargs=True
      )
      module.register_forward_hook(self.forward_post_hook)  # type: ignore[arg-type]
  
      if "device_id" in kwargs:
          # replicate() supports a small usability enhancement where
          # user can pass in device_id as a Union[int, torch.device] even for
          # CPU devices so users don't have to change code for CPU/GPU runs.
          # We derive the right device_ids to feed into DDP to support this.
          if kwargs["device_id"] is not None:
              device_id = kwargs["device_id"]
              # Convert to device_ids that DDP expects.
              if (
                  isinstance(device_id, torch.device)
                  and device_id.type == "cpu"
              ):
                  # CPU modules receive device_ids None
                  kwargs["device_ids"] = None
              else:
                  # GPU modules expect device_ids=[cuda_device]
                  kwargs["device_ids"] = [device_id]
          else:
              kwargs["device_ids"] = None
          kwargs.pop("device_id")
  
      self._ddp = DistributedDataParallel(self._param_list, **kwargs)
      # Weakref to the DDP instance is currently only used for testing.
      replicate.state(self.module)._ddp_weakref = weakref.ref(self._ddp)
  
file_path: pytorch/torch/distributed/_composable/replicate.py
imports: import weakref
  from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
  import torch
  import torch.nn
  from torch.distributed._composable_state import _State
  from torch.nn.parallel import DistributedDataParallel
  from .contract import _get_registry, contract
label_desc: Use image node_id 3 for calling the _ReplicateState obj's underlying member method code with example usage: obj.init(module, ignored_modules) without return types