Confusion on Zero Redundancy Optimizer, seems the parameter shape is CHANGED and REPORT ERROR

[yxding95]

I'm training a model which consists torch.nn.Embedding layer, but when I apply the Zero Redundancy Optimizer, the weight's shape would be changed to 1-dim (torch.Size([50272, 2048]) -> torch.Size([102957056])), and this raises a "RuntimeError: 'weight' must be 2-D" when the embedding is indexed.
Is there anyone know how to solve this problem? Is there any mistake I made? The config.py are just copied from language/gpt/gpt2_configs/gpt2_zero3.py

[ver217]

What is your model?