train - save - restart notebook - load - train - save raises an error

[radekosmulski]

I documented the behavior in this example notebook.

[sgugger]

The problem is in your custom splitter. The last group it returns is not a list but a single parameter. Changing it to:

def splitter(model):
    return params(model)[:150], params(model)[150:-1], params(model)[-1]

makes the bug disappear for me. I'll add some checks in optimizer in any case.