Data Parallel loses ParameterDict

lorr1: <p>Hey, is it expected for ParameterDicts to become empty in DataParallel.</p>
<p>Minimal Example:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.mid = nn.ParameterDict()
        self.mid["key1"] = nn.Parameter(torch.tensor(2.0))
        self.mid["key2"] = nn.Parameter(torch.tensor(3.0))
    def forward(self, x):
        print(self.mid)
        return self.mid["key1"](x) + self.mid["key2"]


print("Devices", torch.cuda.device_count())
mod = MyModule()
mod = nn.DataParallel(mod)
mod = mod.to(0)

input = torch.randn(2, 100, device=0)
out = mod(input)
</code></pre>
<p>Output</p>
<pre><code class="lang-auto">ParameterDict()
ParameterDict()
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 161, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py", line 171, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
    output.reraise()
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/_utils.py", line 428, in reraise
    raise self.exc_type(msg)
KeyError: Caught KeyError in replica 0 on device 0.
Original Traceback (most recent call last):
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "&lt;stdin&gt;", line 9, in forward
  File "/dfs/scratch0/lorr1/py3.8/lib/python3.8/site-packages/torch/nn/modules/container.py", line 542, in __getitem__
    return self._parameters[key]
KeyError: 'key1'
</code></pre>

InnovArul: <aside class="quote no-group quote-modified" data-username="lorr1" data-post="1" data-topic="107227">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lorr1/48/32764_2.png" class="avatar"> lorr1:</div>
<blockquote>
<p><code>return self.mid["key1"](x) + self.min["key2"]</code></p>
</blockquote>
</aside>
<p>You had mis-spelled <code>self.mid</code> as <code>self.min</code>.</p>

lorr1: <p>Thanks for catching that. Even with the spelling error corrected, the error about <code>key1</code> remains. If you see in the print outs, the <code>ParameterDict()</code> is empty.</p>

InnovArul: <pre><code class="lang-auto">/home/arul/miniconda3/envs/pth1.7/lib/python3.7/
site-packages/torch/nn/modules/container.py:645: 
UserWarning: nn.ParameterDict is being used with DataParallel but this is not supported. 
This dict will appear empty for the models replicated on each GPU except the original one.
</code></pre>
<p>It seems <code>nn.ParameterDict</code> is not supported with DataParallel.</p>