code: def parallel_apply(
      modules: Sequence[Module],
      inputs: Sequence[Any],
      kwargs_tup: Optional[Sequence[Dict[str, Any]]] = None,
      devices: Optional[
          Sequence[Optional[Union[int, torch.device]]]
      ] = None,
  ) -> List[Any]:
      r"""Apply each `module` in :attr:`modules` in parallel on each of :attr:`devices`.
  
      Args:
          modules (Module): modules to be parallelized
          inputs (tensor): inputs to the modules
          devices (list of int or torch.device): CUDA devices
  
      :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
      :attr:`devices` (if given) should all have same length. Moreover, each
      element of :attr:`inputs` can either be a single object as the only argument
      to a module, or a collection of positional arguments.
      """
      assert len(modules) == len(
          inputs
      ), f"The number of modules {len(modules)} is not equal to the number of inputs {len(inputs)}"
      if kwargs_tup is not None:
          assert len(modules) == len(kwargs_tup)
      else:
          kwargs_tup = (cast(Dict[str, Any], {}),) * len(modules)
      if devices is not None:
          assert len(modules) == len(devices)
      else:
          devices = [None] * len(modules)
      devices = [_get_device_index(x, True) for x in devices]
      streams = [torch.cuda.current_stream(x) for x in devices]
      lock = threading.Lock()
      results = {}
      grad_enabled, autocast_enabled = (
          torch.is_grad_enabled(),
          torch.is_autocast_enabled(),
      )
  
      def _worker(
          i: int,
          module: Module,
          input: Any,
          kwargs: Dict[str, Any],
          device: Optional[Union[int, torch.device]] = None,
          stream: Optional[torch.cuda.Stream] = None,
      ) -> None:
          torch.set_grad_enabled(grad_enabled)
          if device is None:
              t = get_a_var(input)
              if t is None:
                  with lock:
                      results[i] = ExceptionWrapper(
                          where=f"in replica {i}, no device was provided and no tensor input was found; "
                          "device cannot be resolved"
                      )
                  return
              device = t.get_device()
          if stream is None:
              stream = torch.cuda.current_stream(device)
          try:
              with torch.cuda.device(device), torch.cuda.stream(
                  stream
              ), autocast(enabled=autocast_enabled):
                  # this also avoids accidental slicing of `input` if it is a Tensor
                  if not isinstance(input, (list, tuple)):
                      input = (input,)
                  output = module(*input, **kwargs)
              with lock:
                  results[i] = output
          except Exception:
              with lock:
                  results[i] = ExceptionWrapper(
                      where=f"in replica {i} on device {device}"
                  )
  
      if len(modules) > 1:
          threads = [
              threading.Thread(
                  target=_worker,
                  args=(i, module, input, kwargs, device, stream),
              )
              for i, (
                  module,
                  input,
                  kwargs,
                  device,
                  stream,
              ) in enumerate(
                  zip(modules, inputs, kwargs_tup, devices, streams)
              )
          ]
  
          for thread in threads:
              thread.start()
          for thread in threads:
              thread.join()
      else:
          _worker(
              0,
              modules[0],
              inputs[0],
              kwargs_tup[0],
              devices[0],
              streams[0],
          )
  
      outputs = []
      for i in range(len(inputs)):
          output = results[i]
          if isinstance(output, ExceptionWrapper):
              output.reraise()
          outputs.append(output)
      return outputs
  
file_path: pytorch/torch/nn/parallel/parallel_apply.py
imports: import threading
  import torch
  from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
  from ..modules import Module
  from torch.cuda._utils import _get_device_index
  from torch.cuda.amp import autocast
  from torch._utils import ExceptionWrapper
label_desc: Use image node_id 2 for calling a global function with example usage: parallel_apply(modules, inputs, kwargs_tup, devices) and returns: outputs