shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1) ValueError: not enough values to unpack (expected 6, got 4)

[LianShuaiLong]

Discussed in #9710

^{Originally posted by LianShuaiLong October 18, 2024}
To save GPU memory during training, I pre-extracted the text embeddings. During the training process, I did not load the text encoders and directly used the embeddings for training. However, I encountered the following error:

shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
ValueError: not enough values to unpack (expected 6, got 4)

Here is my feature extraction script and the modified training script. Could someone kindly help me identify where the issue might be?