Output switches to CPU after passing though ModuleList blocks

username: saandeep_aathreya
cooked: <p>I am training a model using DistributedDataParallel(code snippet below). My model is initialized using nn.ModuleList(). But once the input which is on GPU passes through one of the blocks in nn.ModuleList, it switches to CPU mode.</p>
<pre><code class="lang-auto">if use_cuda and torch.cuda.device_count()&gt;1:
    model = model.to(rank)
    model = DistributedDataParallel(model, device_ids=[rank])
</code></pre>
<p>Please refer to <em>forward</em> method of <strong>Block</strong> class which creates a list of <strong>SimpleNet</strong> class. Please let me know if you require additional piece of code.</p>
<pre><code class="lang-auto">

class SimpleNet(nn.Module):
  def __init__(self, inp, parity):
    super(SimpleNet, self).__init__()
    self.net = nn.Sequential(
      nn.Linear(inp//2, 256),
      nn.LeakyReLU(True),
      nn.Linear(256, 256),
      nn.LeakyReLU(True),
      nn.Linear(256, inp//2),
      nn.Sigmoid(),
      nn.BatchNorm1d(392)
    )
    self.inp = inp
    self.parity = parity

  def forward(self, x):
    z = torch.zeros(x.size())
    x0, x1 = x[:, ::2], x[:, 1::2]
    if self.parity % 2:
      x0, x1 = x1, x0 
    # print("X: ", x0[0][0].detach(), x1[0][0].detach())
    z1 = x1
    log_s = self.net(x1)
    # print(x.size(), x1.size(), log_s.size())
    t = self.net(x1)
    s = torch.exp(log_s)
    z0 = (s * x0) + t
    # print("Z: ", z0[0][0].detach(), z1[0][0].detach())
    if self.parity%2:
      z0, z1 = z1, z0
    z[:, ::2] = z0
    z[:, 1::2] = z1
    logdet = torch.sum(torch.log(s), dim = 1)
    return z, logdet
  
  def reverse(self, z):
    x = torch.zeros(z.size())
    z0, z1 = z[:, ::2], z[:, 1::2]
    if self.parity%2:
      z0, z1 = z1, z0
    # print("Z: ", z0[0][0].detach(), z1[0][0].detach())
    x1 = z1
    log_s = self.net(z1)
    t = self.net(z1)
    s = torch.exp(log_s)
    x0 = (z0 - t)/s
    # print("X: ", x0[0][0].detach(), x1[0][0].detach())
    if self.parity%2:
      x0, x1 = x1, x0
    x[:, ::2] = x0
    x[:, 1::2] = x1
    return x
  

class Block(nn.Module):
  def __init__(self, inp, n_blocks):
    super(Block, self).__init__()
    parity = 0
    blocks = []
    for _ in range(n_blocks):
      blocks.append(SimpleNet(inp, parity))
      parity += 1
    self.blocks = nn.ModuleList(blocks) 
  
  def forward(self, x):
    logdet = 0
    out = x
    xs = [out]
    # print("*"*20, "FORWARD", "*"*30)
    for block in self.blocks:
      print("device in block: ", out.is_cuda) # True
      out, det = block(out)
      print("device in block: ", out.is_cuda) # False
      logdet += det
      xs.append(out)
    return out, logdet

  def reverse(self, z):
    # print("*"*20, "REVERSE", "*"*30)
    out = z
    for block in self.blocks[::-1]:
      out = block.reverse(out)
    return out
</code></pre>

username: osalpekar
cooked: <p>In the <code>SimpleNet</code> forward() function, the zeros tensor <code>z</code> is being created on CPU, which is why the <code>out</code> tensor returned by the forward function of that block returns <code>out.is_cuda = False</code>. You must either place the <code>z</code> tensor explicitly on the correct rank or make it an <code>nn.Parameter</code> or similar type that will be moved to GPU when the entire module is placed on GPU.</p>

username: saandeep_aathreya
cooked: <p>Thank you for your response. I was able to make it work using the <code>nn.Parameter</code> and <code>register_parameter</code> but I noticed that if I just initialize the <code>z</code> tensor like <code>z = torch.zeros_like(x)</code> instead of <code>torch.zeros(x.size())</code>, it is automatically loaded to the same device as <code>x</code>.</p>