Skip to content

Commit

Permalink
[mis][ci/test] fix flaky test in test_sharded_state_loader.py (vllm-p…
Browse files Browse the repository at this point in the history
…roject#5361)

[mis][ci/test] fix flaky test in tests/test_sharded_state_loader.py (vllm-project#5361)
  • Loading branch information
youkaichao authored and jimpang committed Jun 27, 2024
1 parent 8514bbd commit 8bc060c
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion tests/test_sharded_state_loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ def test_filter_subtensors():
filtered_state_dict = ShardedStateLoader._filter_subtensors(state_dict)
assert tuple(filtered_state_dict.keys()) == ("a", "b", "c")
for key, tensor in filtered_state_dict.items():
assert tensor.equal(state_dict[key])
# NOTE: don't use `euqal` here, as the tensor might contain NaNs
assert tensor is state_dict[key]


@pytest.fixture(scope="module")
Expand Down

0 comments on commit 8bc060c

Please sign in to comment.