ENAS bug IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

[marsggbo]

In ENAS example, when lstm_num_layers=1, the code can run successfully. However, when I set lstm_num_layers=2, then it raises the following error: IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

[marsggbo]

I solve this problem by modifying the output shape in the forward function, as below:

class StackedLSTMCell(nn.Module):
    def __init__(self, layers, size, bias):
        super().__init__()
        self.lstm_num_layers = layers
        self.lstm_modules = nn.ModuleList([nn.LSTMCell(size, size, bias=bias)
                                           for _ in range(self.lstm_num_layers)])

    def forward(self, inputs, hidden):
        prev_c, prev_h = hidden
        next_c, next_h = [], []
        for i, m in enumerate(self.lstm_modules):
            curr_c, curr_h = m(inputs, (prev_c[i], prev_h[i]))
            next_c.append(curr_c)
            next_h.append(curr_h)
            inputs = curr_h[-1].view(1,-1) # here
        return next_c, next_h

[ultmaster]

Great! Since you've solved the problem, would you like to submit a pull request? Thanks.