pytorch dataparallell unexpected behaviour

rnntest.py

import numpy as np 
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F

print('torch.__version__',torch.__version__)
print('GPUs:')
for k in range(torch.cuda.device_count()):
  print('\t',torch.cuda.get_device_properties(k)) 
# torch.set_default_tensor_type('torch.cuda.FloatTensor')
n_gpu = torch.cuda.device_count()

class RNN(nn.Module):
    """
        GRU with trainable h0
    """
    def __init__(self, n_input, n_output, n_hidden=10, n_layers=1):
        super(RNN, self).__init__()
        self.linear_in = nn.Linear(n_input, n_hidden)

        self.gru = nn.GRU(n_hidden, n_hidden, n_layers)        
        self.linear_out = nn.Linear(n_hidden, n_output,bias=False)
        self.gru.flatten_parameters()

    def forward(self, x):
        x = self.linear_in(x).tanh()
        x, hidden = self.gru(x)
        return self.linear_out(x),hidden

n_timesteps = 1000
n_batch = 500
n_features = 100

x = torch.randn([n_timesteps,n_batch,n_features]).cuda()
# rnn = nn.GRU(input_size=n_features, hidden_size=1, num_layers=1).cuda()
rnn = RNN(n_input = n_features,n_output =1, n_hidden = 10).cuda()
rnn = nn.DataParallel(rnn)

######### PLOT
plt.imshow(y.T)
plt.xlabel('time')
plt.ylabel('batch')
plt.savefig('./result/timeline_'+str(n_gpu)+ '_gpus.png')
plt.close()

fig, ax = plt.subplots(ncols=1, nrows=2, sharex=True, figsize=(6, 4))

ax[0].imshow(y.T, interpolation='none')
ax[0].set_ylabel('batch')
ax[0].set_xlabel('time')

ax[1].plot(np.nanmean(y.T, axis=0), lw=0.5,c='black', drawstyle='steps-post')
ax[1].set_title('mean/timestep')
plt.savefig('./result/timeline_agg_'+str(n_gpu)+ '_gpus.png')
plt.close()
print('finished')

x = torch.randn([n_timesteps,n_batch,n_features]).cuda()
# rnn = nn.GRU(input_size=n_features, hidden_size=1, num_layers=1).cuda()
rnn = RNN(n_input = n_features,n_output =1, n_hidden = 10).cuda()
rnn = nn.DataParallel(rnn,dim=1) # This is the fix

######### PLOT
plt.imshow(y.T)
plt.xlabel('time')
plt.ylabel('batch')
plt.savefig('./result/timeline_dim1'+str(n_gpu)+ '_gpus.png')
plt.close()

fig, ax = plt.subplots(ncols=1, nrows=2, sharex=True, figsize=(6, 4))

ax[0].imshow(y.T, interpolation='none')
ax[0].set_ylabel('batch')
ax[0].set_xlabel('time')

ax[1].plot(np.nanmean(y.T, axis=0), lw=0.5,c='black', drawstyle='steps-post')
ax[1].set_title('mean/timestep')
plt.savefig('./result/timeline_agg_dim1'+str(n_gpu)+ '_gpus.png')
plt.close()
print('finished')