I am new to pytorch, here is my question. I have implemented ConvLSTM on pytorch by I could not find a way to initialize the hidden states before unrolling the ConvLSTM. This is the ConvLSTM cell and layer.
import torch
import torch.nn as nn
from torch.autograd import Variable


class ConvLSTMCell(nn.Module):
    def __init__(self, input_channels, hidden_channels, kernel_size):
        super(ConvLSTMCell, self).__init__()

        #assert hidden_channels % 2 == 0

        self.input_channels = input_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.num_features = 4

        self.padding = int((kernel_size - 1) / 2)

        self.Wxi = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Whi = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)
        self.Wxf = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Whf = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)
        self.Wxc = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Whc = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)
        self.Wxo = nn.Conv2d(self.input_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=True)
        self.Who = nn.Conv2d(self.hidden_channels, self.hidden_channels, self.kernel_size, 1, self.padding, bias=False)

        self.Wci = None
        self.Wcf = None
        self.Wco = None

    def forward(self, x, h, c):
        ci = torch.sigmoid(self.Wxi(x) + self.Whi(h) + c * self.Wci)
        cf = torch.sigmoid(self.Wxf(x) + self.Whf(h) + c * self.Wcf)
        cc = cf * c + ci * torch.tanh(self.Wxc(x) + self.Whc(h))
        co = torch.sigmoid(self.Wxo(x) + self.Who(h) + cc * self.Wco)
        ch = co * torch.tanh(cc)
        return ch, cc

    def init_hidden(self, batch_size, hidden, shape):
        if self.Wci is None:
            self.Wci = nn.Parameter(torch.zeros(1, hidden, shape[0], shape[1])).cuda()
            self.Wcf = nn.Parameter(torch.zeros(1, hidden, shape[0], shape[1])).cuda()
            self.Wco = nn.Parameter(torch.zeros(1, hidden, shape[0], shape[1])).cuda()
        else:
            assert shape[0] == self.Wci.size()[2], 'Input Height Mismatched!'
            assert shape[1] == self.Wci.size()[3], 'Input Width Mismatched!'
        return (nn.Parameter(torch.zeros(batch_size, hidden, shape[0], shape[1])).cuda(),
                nn.Parameter(torch.zeros(batch_size, hidden, shape[0], shape[1])).cuda())

class ConvLSTM_s(nn.Module):
    # input_channels corresponds to the first input feature map
    # hidden state is a list of succeeding lstm layers.
    def __init__(self, input_channels, hidden_channels, kernel_size, step=1, effective_step=[1]):
        super(ConvLSTM_s, self).__init__()
        self.input_channels = [input_channels] + hidden_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.num_layers = len(hidden_channels)
        self.step = step
        self.effective_step = effective_step
        self._all_layers = []
        for i in range(self.num_layers):
            name = 'cell{}'.format(i)
            cell = ConvLSTMCell(self.input_channels[i], self.hidden_channels[i], self.kernel_size)
            setattr(self, name, cell)
            self._all_layers.append(cell)

    def forward(self, inputx, states=None):
        internal_state = []
        outputs = []
        for step in range(self.step):

            x = inputx[:,step,:,:,:] #slicing the time window
            #print(x.shape)
            for i in range(self.num_layers):
                # all cells are initialized in the first step
                name = 'cell{}'.format(i)
                if step == 0:
                    bsize, _, height, width = x.size()
                    if states == None:
                    (h, c) = getattr(self, name).init_hidden(batch_size=bsize, hidden=self.hidden_channels[i],                                          shape=(height, width))
                    else:
                    (h,c) = states
                    
                    internal_state.append((h, c))

                # do forward
                (h, c) = internal_state[i]
                x, new_c = getattr(self, name)(x, h, c)
                internal_state[i] = (x, new_c)
            # only record effective steps
            #if step in self.effective_step:
            if True:
                outputs.append(x)
        assert len(outputs) == inputx.shape[1]
        outputs = torch.stack(outputs,dim=1)
        return (outputs, internal_state[-1])

if __name__ == '__main__':

    inputx = Variable(torch.randn(1, 3, 4, 64, 32)).cuda()

    h1 = (torch.zeros(1, 2, 64, 32)).cuda()
    h2 = (torch.zeros(1, 2, 64, 32)).cuda()

    target = Variable(torch.randn(1, 3, 2, 64, 32)).double().cuda()

    convlstms = ConvLSTM_s(input_channels=4, hidden_channels=[2], kernel_size=3, step=3,
                        effective_step=[3]).cuda()
    loss_fn = torch.nn.MSELoss()




    output = convlstms(inputx,states=(h1,h2))

And the error is
     32     def forward(self, x, h, c):
---> 33         ci = torch.sigmoid(self.Wxi(x) + self.Whi(h) + c * self.Wci)
     34         cf = torch.sigmoid(self.Wxf(x) + self.Whf(h) + c * self.Wcf)
     35         cc = cf * c + ci * torch.tanh(self.Wxc(x) + self.Whc(h))

TypeError: mul(): argument 'other' (position 1) must be Tensor, not NoneType

It stems from the operations in the ConvLSTM cell