if you manually implement DataParallel yourself, this is possible.
We implemented DataParallel using collectives such as broadcast, scatter, gather.
See this function for reference (and if you want to copy it and modify yourself):

  
      github.com
  
  
    pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L76-L106 12



def parallel_apply(self, replicas, inputs, kwargs):
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])


def gather(self, outputs, output_device):
    return gather(outputs, output_device, dim=self.dim)




def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None):
r"""Evaluates module(input) in parallel across the GPUs given in device_ids.


This is the functional version of the DataParallel module.


Args:
    module: the module to evaluate in parallel
    inputs: inputs to the module
    device_ids: GPU ids on which to replicate module
    output_device: GPU location of the output  Use -1 to indicate the CPU.
        (default: device_ids[0])
Returns:


  This file has been truncated. show original