DiT forward pass assertion fails

[HuronExplodium]

Following install guide + flash attention:

python3 ./demos/cli.py --model_dir weights/ --prompt "a greyhound running"

results in:

  File "/home/user/video/models/src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py", line 703, in forward
    assert self.num_heads % cp_size == 0
AssertionError

[ved-genmo]

Is this running on one GPU or multiple GPUs?

[HuronExplodium]

Initially tried with 7x3090 but then restricted to 2x with CUDA_VISIBLE_DEVICES and got the same thing.