title: How to parallelize RNN function in Pytorch with DataParallel
body: <p>Here's an RNN model to run character based language generation:</p>
  
  <pre><code>class RNN(nn.Module):
      def __init__(self, input_size, hidden_size, output_size, n_layers):
          super(RNN, self).__init__()
          self.input_size = input_size
          self.hidden_size = hidden_size
          self.output_size = output_size
          self.n_layers = n_layers 
  
          self.encoder = nn.Embedding(input_size, hidden_size)
          self.GRU = nn.GRU(hidden_size, hidden_size, n_layers, batch_first=True) 
          self.decoder = nn.Linear(hidden_size, output_size)
  
  
      def forward(self, input, batch_size):
          self.init_hidden(batch_size)
          input = self.encoder(input)
          output, self.hidden = self.GRU(input, self.hidden) 
          output = self.decoder(output.view(batch_size, self.hidden_size)) 
          return output
  
      def init_hidden(self, batch_size):
          self.hidden = Variable(torch.randn(self.n_layers, batch_size, self.hidden_size).cuda())
  </code></pre>
  
  <p>I instantiate the model using DataParallel, to split the batch of inputs across my 4 GPUs:</p>
  
  <pre><code>net = torch.nn.DataParallel(RNN(n_chars, hidden_size, n_chars, n_layers)).cuda()
  </code></pre>
  
  <p>Here's the <a href="https://gist.github.com/michaelklachko/26565dfac747d628961901d3dd84bd71" rel="nofollow noreferrer">full code</a>.</p>
  
  <p>Unfortunately, DataParallel requires the inputs to have batch_size as the first dimension, but GRU function expects hidden tensor to have batch_size as second dimension:</p>
  
  <pre><code>output, self.hidden = self.GRU(input, self.hidden)
  </code></pre>
  
  <p>The code as is throws the following error (note the printouts showing that encoder is correctly executed on 4 GPUs):</p>
  
  <pre><code>...
  forward function: encoding input of shape: (16L, 1L)
  forward function: encoding input of shape: (16L, 1L)
  forward function: encoding input of shape: (16L,
  forward function: encoding input of shape:
  
  forward function: GRU processing input of shape:
  1L)
   ( (16L, 16L1L, 1L), 100L)
  forward function: GRU processing input of shape:
   (16L, 1L,
  forward function: GRU processing input of shape:100L)
   (16L
  forward function: GRU processing input of shape:, 1L, 100L) (
  16L, 1L, 100L)
  
  Traceback (most recent call last):
    File "gru2.py", line 166, in &lt;module&gt;
      output = net(c, batch_size)
    File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/modules/module.py", line 206, in __call__
      result = self.forward(*input, **kwargs)
    File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/parallel/data_parallel.py", line 61, in forward
      outputs = self.parallel_apply(replicas, inputs, kwargs)
    File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/parallel/data_parallel.py", line 71, in parallel_apply
      return parallel_apply(replicas, inputs, kwargs)
    File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/parallel/parallel_apply.py", line 45, in parallel_apply
      raise output
  RuntimeError: Expected hidden size (2, 16L, 100), got (2L, 64L, 100L)
  </code></pre>
  
  <p>Here the model has 2 layers, batch_size=64, and hidden_size = 100. </p>
  
  <p>How do I parallelize the GRU operation in the forward function? </p>