How are pretrained weights loaded in nn.LazyLinear layers

Hello in a project of mine I am trying to build a linear layer for which the input dim may change depending upon the dataset used.
So I am using a nn.lazylinear.
However after obtaining the stat_dict after training whenever I try to load the model using the state_dict I get issues with the keys :
untimeError: Error(s) in loading state_dict for SeamModule:

    Unexpected key(s) in state_dict: "model.enc_to_dec.weight", "model.enc_to_dec.bias". 

How do i properly load the trained Lazylinear weights ?

I’m unable to reproduce the issue by trying to load a state_dict into an uninitialized Lazy* model. Could you post a minimal code snippet to reproduce the issue, please?

Also, as a workaround you might need to perform the dry-run before loading the state_dict but it’s unclear to my if that’s indeed necessary as I don’t see any errors.

Yes, i tried training a model again from scratch using nn.LazyLinear. It seems to work fine when I later load checkpoints.
I guess the error was coming up from a different part of the code.
Thanks for the prompt reply though !