Skip to content

Commit 8eca187

Browse files
committed
update test torch version
Signed-off-by: Ye Yu <yeyu@nvidia.com>
1 parent bc4eeb4 commit 8eca187

1 file changed

Lines changed: 2 additions & 5 deletions

File tree

tests/examples/speculative_decoding/test_eagle.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,8 @@ def test_llama_eagle3(tiny_llama_path, tiny_daring_anteater_path, tmp_path, eagl
3737
available_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
3838
if cp_size == 2 and available_gpus < 2:
3939
pytest.skip("cp_size=2 requires at least 2 GPUs, but only {} found.".format(available_gpus))
40-
if cp_size == 2 and not (
41-
Version(torch.__version__) > Version("2.7.0")
42-
and Version(torch.__version__) < Version("2.9.0")
43-
):
44-
pytest.skip("cp_size=2 requires torch 2.8.0")
40+
if cp_size == 2 and not Version(torch.__version__) >= Version("2.10.0"):
41+
pytest.skip("cp_size=2 requires torch 2.10.0")
4542
# Create an ultra-tiny EAGLE config for testing to reduce memory usage
4643
tiny_eagle_config = {
4744
"max_position_embeddings": 128,

0 commit comments

Comments
 (0)