Problem with Skorch and RNN with LSTM cells

I am testing Skorch with LSTM cells for a regression problem. Here is the code

import numpy as np
from sklearn.datasets import make_regression
from skorch import NeuralNetRegressor
import unittest
from util.pythorch_util import MyLSTM

        X_regr, y_regr = make_regression(1000, 20, n_informative=10, random_state=0)
        X_regr = X_regr.astype(np.float32)
        y_regr = y_regr.astype(np.float32) / 100
        y_regr = y_regr.reshape(-1, 1)
        ni = 20
        no = 1
        nh = 10
        nlayers = 3
        net_regr = NeuralNetRegressor(
            module=MyLSTM,
            module__ni=ni,
            module__no=no,
            module__nh=nh,
            module__nlayers=nlayers,
            max_epochs=20,
            lr=0.1,
            #     device='cuda',  # uncomment this to train with CUDA
        )

        net_regr.fit(X_regr,y_regr)

where MYLSTM is

class MyLSTM(nn.Module):
    def __init__(self, ni=6, no=3, nh=10, nlayers=1):
        super(MyLSTM, self).__init__()

        self.ni = ni
        self.no = no
        self.nh = nh
        self.nlayers = nlayers

        self.lstms = nn.ModuleList(
            [nn.LSTMCell(self.ni, self.nh)] + [nn.LSTMCell(self.nh, self.nh) for i in range(nlayers - 1)])
        self.out = nn.Linear(self.nh, self.no)
        self.do = nn.Dropout(p=0.2)
        self.actfn = nn.Tanh()
        self.device = torch.device('cpu')
        self.dtype = torch.float

    # description of the whole block
    def forward(self, x, h0=None, train=False):
        hs = x  # initiate hidden state
        if h0 is None:
            h = torch.zeros(hs.shape[0], self.nh, device=device)
            c = torch.zeros(hs.shape[0], self.nh, device=device)
        else:
            (h, c) = h0

        # LSTM cells
        for i in range(self.nlayers):
            h, c = self.lstms[i](hs, (h, c))
            if train:
                hs = self.do(h)
            else:
                hs = h
        y = self.out(hs)
        return y, (h, c)

when I try to run I get this error:

Error
Traceback (most recent call last):
  File "C:\Users\morabru01\AppData\Local\Programs\Python\Python36\lib\unittest\case.py", line 59, in testPartExecutor
    yield
  File "C:\Users\morabru01\AppData\Local\Programs\Python\Python36\lib\unittest\case.py", line 605, in run
    testMethod()
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\testbenches\test_LSTM_skorch.py", line 30, in test_toy
    net_regr.fit(X_regr,y_regr)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\regressor.py", line 91, in fit
    return super(NeuralNetRegressor, self).fit(X, y, **fit_params)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 854, in fit
    self.partial_fit(X, y, **fit_params)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 813, in partial_fit
    self.fit_loop(X, y, **fit_params)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 727, in fit_loop
    step_fn=self.train_step, **fit_params)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 763, in run_single_epoch
    step = step_fn(Xi, yi, **fit_params)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 659, in train_step
    self.optimizer_.step(step_fn)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\torch\optim\sgd.py", line 80, in step
    loss = closure()
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 656, in step_fn
    step = self.train_step_single(Xi, yi, **fit_params)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 599, in train_step_single
    loss = self.get_loss(y_pred, yi, X=Xi, training=True)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\skorch\net.py", line 1102, in get_loss
    return self.criterion_(y_pred, y_true)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\torch\nn\modules\module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\torch\nn\modules\loss.py", line 431, in forward
    return F.mse_loss(input, target, reduction=self.reduction)
  File "C:\Users\morabru01\Desktop\WORKSPACE\Learning\venv\lib\site-packages\torch\nn\functional.py", line 2203, in mse_loss
    if not (target.size() == input.size()):
AttributeError: 'tuple' object has no attribute 'size'

I am using Python 3.6, torch 1.4.0 and skorch 0.8.

EDIT: I guess that the problem comes from the fact that lstm takes a tuple (h,c) with the hidden states… is there a way to solve this?

Solved here