@staticmethod
    def forward(ctx, input, weights, bias, old_h, old_cell):
        outputs = lltm.forward(input, weights, bias, old_h, old_cell)
        new_h, new_cell = outputs[:2]
#        variables = outputs[1:] + [weights, old_cell] # Doesn't work
        variables = outputs[2:] + [weights, old_cell]
        ctx.save_for_backward(*variables)

        return new_h, new_cell

Modify to this
variables = outputs[2:] + [weights]