[mis][ci/test] fix flaky test in test_sharded_state_loader.py (#5361)

[mis][ci/test] fix flaky test in tests/test_sharded_state_loader.py (#5361)
vllm-project · Jun 9, 2024 · 5d7e3d0 · 5d7e3d0
1 parent 0373e18
commit 5d7e3d0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_sharded_state_loader.py b/tests/test_sharded_state_loader.py
@@ -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")