Fix optimizer reloading from checkpoint

[awni]

This is kind of a weird problem.

Basically if you have a model like:

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(2, 2)
        self.drop = nn.Dropout(p=0.5)
        self.l2 = nn.Linear(2, 2)

It has a parameters() like {"l1": ..., "drop": {}, "l2": ...}. The optimizer state is setup to match that. When serializing the optimizer state it gets flattened which removes the empty drop value from the tree. That in turn causes issues when deserializing the optimizer because its' state no longer matches that of parameters().

I think there are two possible fixes here, I implemented the first, but I'm open to trying the second:

• Delay completing the initialization of the optimizer until we have the first parameter state so we can match it.
• Change Module.parameters() and Module.trainable_parameters() to not include empty values.

This closes #1328

[angeloskath]

Other than the one comment I think this looks good.

I am not too happy that the tree walking logic starts crawling into everything but so far it is not too annoying.