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

[wangmiaowei]

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