Model architecture is different for DiT-B/2 Scratch ckpt (omega_embedder) #6
Description
Hello,
When I tried loading the provided checkpoint DiT-B-GF-scratch.pt, I encountered the following error:
[rank0]: File "/opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2215, in load_state_dict
[rank0]: raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
[rank0]: RuntimeError: Error(s) in loading state_dict for DiT:
[rank0]: Missing key(s) in state_dict: "t_embedder.cond_proj.weight".
[rank0]: Unexpected key(s) in state_dict: "omega_embedder.mlp.0.weight", "omega_embedder.mlp.0.bias", "omega_embedder.mlp.2.weight", "omega_embedder.mlp.2.bias".
It seems that the checkpoint and the current model code have slight differences in the embedder structure (perhaps due to version differences). While this is a minor issue, it would be very helpful for reproducibility if the author could provide the architecture of the omega_embedder.
Thank you for your time and for sharing this work!