test: Fix test_sampling.py on Spark#2042

yzh119 merged 1 commit intoflashinfer-ai:mainflashinfer-ai/flashinfer:mainfrom

bkryu:spark_test_sampling_fixbkryu/flashinfer:spark_test_sampling_fixCopy head branch name to clipboard