[BUG] Diffusion Gemma run benchmark error #45

Required prerequisites

• I have read the project documentation https://github.com/SJTU-DENG-Lab/Diffulex#readme.
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version/commit of Diffulex are you using?

184ba97

System information

python:3.11.2 (main, Apr 28 2025, 14:11:48) [GCC 12.2.0] linux
torch:2.11.0+cu130

Problem description

Hi.When I run the benchmark :CUDA_VISIBLE_DEVICES=0 python -m diffulex_bench.main
--config diffulex_bench/configs/diffusion_gemma_gsm8k.yml
--model-path /path/to/diffusiongemma-26B-A4B-it
--dataset-limit 10
It reports:

Reproducible example code

The Python snippets:

CUDA_VISIBLE_DEVICES=0 python -m diffulex_bench.main \
  --config diffulex_bench/configs/diffusion_gemma_gsm8k.yml \
  --model-path /path/to/diffusiongemma-26B-A4B-it \
  --dataset-limit 10

Traceback

Expected behavior

No response

Additional context

No response