language model load from checkpoint error

[omerarshad]

🐛 Bug

Saving aggregated checpoint for language modeling transformer gives error

RuntimeError: Error(s) in loading state_dict for LanguageModelingTransformer:
	Missing key(s) in state_dict: "model.lm_head.weight". 

To Reproduce

from pytorch_lightning.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict

convert_zero_checkpoint_to_fp32_state_dict(
    "./recreate_model/epoch=0-step=363.ckpt/",
    "./recreate_model/pytorch_model.bin"
 )

# Load best model from aggregated checkpoint file
best_model = LanguageModelingTransformer.load_from_checkpoint(
    "./recreate_model/pytorch_model.bin"
)

[Borda]

could you please share the full trace? 🐰
seems to be a duplicate of #273 (comment) so lets keep only one 🦦