Bug: Error: only integer tensors of a single element can be converted to an index

Reproduce the failure with the torchrun command in the issue, then start at generate.py:172 and follow the call into pipeline/inference_pipeline.py:602. Inspect wan/utils/fm_solvers_unipc.py around convert_model_output at line 323 and step at line 697, using the traceback to identify the incompatible tensor operation. Done means the command completes past this error and produces the requested output.

My num_gpus is 6:
and test with the command:

torchrun --nproc_per_node=$NUM_GPUS generate.py --size 704*1280 --dit_fsdp --t5_fsdp --ckpt_dir Matrix-Game-3.0 --fa_version 3 --use_int8 --num_iterations 12 --num_inference_steps 3 --image demo_images/001/image.png --prompt "A colorful, animated cityscape with a gas station and various buildings." --save_name test --seed 42 --compile_vae --lightvae_pruning_rate 0.5 --vae_type mg_lightvae --output_dir ./output

ank4]: Traceback (most recent call last):
[rank4]: File "/mnt/group/miaoweiwang/survey/Matrix-Game/Matrix-Game-3/generate.py", line 193, in
[rank4]: generate(args)
[rank4]: File "/mnt/group/miaoweiwang/survey/Matrix-Game/Matrix-Game-3/generate.py", line 172, in generate
[rank4]: pipeline.generate(
[rank4]: File "/mnt/group/miaoweiwang/survey/Matrix-Game/Matrix-Game-3/pipeline/inference_pipeline.py", line 602, in generate
[rank4]: latents = test_scheduler.step(
[rank4]: ^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/mnt/group/miaoweiwang/survey/Matrix-Game/Matrix-Game-3/wan/utils/fm_solvers_unipc.py", line 697, in step
[rank4]: model_output_convert = self.convert_model_output(
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/mnt/group/miaoweiwang/survey/Matrix-Game/Matrix-Game-3/wan/utils/fm_solvers_unipc.py", line 323, in convert_model_output
[rank4]: x0_pred = sample - sigma_t * model_output
[rank4]: ~~~~~~~~^~~~~~~~~~~~~~
[rank4]: TypeError: only integer tensors of a single element can be converted to an index