I have a very simple LSTM model:
class MyLSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super(MyLSTM,self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.LSTM = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.LNN = nn.Linear(hidden_dim, input_dim)

    def forward(self, i, hc):
        o, _ = self.LSTM(i, hc)
        return self.LNN(o)

which I train with the following code:
params = {'batch_size': args.batch,
          'shuffle': False,
          'drop_last': True}
dataloader = DataLoader(train_set, **params)

if cuda:
    model.cuda()

hidden_dim = n_features
model = MyLSTM(n_features, hidden_dim)

optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr)
loss_f = nn.MSELoss()

for epoch in range(args.epochs):
    for i, (seq, truth) in enumerate(dataloader):
        model.zero_grad()
        optimizer.zero_grad()
        h = torch.rand(args.batch * hidden_dim).view(1, args.batch, hidden_dim)
        c = torch.rand(args.batch * hidden_dim).view(1, args.batch, hidden_dim)
        output = model(seq, (h, c))
        loss = loss_f(output, truth)
        loss.backward()
        optimizer.step()

This works fine, but now I would like to use DataParallel to train my model on 2 GPUs.
So I changed the above code to:
if cuda:
    if len(gpus) > 1:
        model = nn.DataParallel(model, device_ids=[0,1])
    model.cuda()

Also, it seems that the sizes of the hidden vectors must be adjusted as well:
h = torch.rand(args.batch // n_gpus * hidden_dim).view(1, args.batch // n_gpus, hidden_dim)
c = torch.rand(args.batch // n_gpus * hidden_dim).view(1, args.batch // n_gpus, hidden_dim)

Now, I face another problem (assuming batch size = 16):
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/loss.py:431: UserWarning: Using a target size (torch.Size([16, 1440, 3969])) that is different to the input size (torch.Size([8, 1440, 3969])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
  return F.mse_loss(input, target, reduction=self.reduction)
Traceback (most recent call last):
  File "py/train.py", line 144, in <module>
    loss = loss_f(output, truth)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/loss.py", line 431, in forward
    return F.mse_loss(input, target, reduction=self.reduction)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 2215, in mse_loss
    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
  File "/usr/local/lib/python3.6/dist-packages/torch/functional.py", line 52, in broadcast_tensors
    return torch._C._VariableFunctions.broadcast_tensors(tensors)
RuntimeError: The size of tensor a (8) must match the size of tensor b (16) at non-singleton dimension 0

After a bit of digging, it seems that, somehow, only 1 replica of my model is used. That is, each replica correctly receives 1/2 of the batch (batch size = 8), but having only 1 replica, the final output is not 16, but only 8, hence the size mismatch error.
This is confirmed by the fact that, if I put one print(o.size()) inside the forward method of the MyLSTM model, I see only 1 torch.Size([8, T, n_features]), while I would expect 2.
Any idea on the possible causes of this?