Hi,
Seems like this question had come up before. Seems like this limitation is intrinsic to DataParallel, e.g. see Yanli_Zhao answer here:
  
    
    
    When I run model on multiple GPUs,register_hook is invalid distributed
  
  
    I want to save gradients of internal variables through register_hook() or retain_grad(). 
When I run model on single GPU, it works. 
But when I run model on multiple GPUs through wrapping model into nn.DataParallel, I find that it doesn’t work. 
Can anyone help me?
  


Also see warnings here
  
      github.com
  
  
    pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L67

Arbitrary positional and keyword inputs are allowed to be passed into
DataParallel but some types are specially handled. tensors will be
**scattered** on dim specified (default 0). tuple, list and dict types will
be shallow copied. The other types will be shared among different threads
and can be corrupted if written to in the model's forward pass.

The parallelized :attr:`module` must have its parameters and buffers on
``device_ids[0]`` before running this :class:`~torch.nn.DataParallel`
module.

.. warning::
    In each forward, :attr:`module` is **replicated** on each device, so any
    updates to the running module in ``forward`` will be lost. For example,
    if :attr:`module` has a counter attribute that is incremented in each
    ``forward``, it will always stay at the initial value because the update
    is done on the replicas which are destroyed after ``forward``. However,
    :class:`~torch.nn.DataParallel` guarantees that the replica on
    ``device[0]`` will have its parameters and buffers sharing storage with
    the base parallelized :attr:`module`. So **in-place** updates to the
    parameters or buffers on ``device[0]`` will be recorded. E.g.,
    :class:`~torch.nn.BatchNorm2d` and :func:`~torch.nn.utils.spectral_norm`