Skip to content

Commit 218ff20

Browse files
committed
[CI] add missing argument
Signed-off-by: Andy Xie <[email protected]>
1 parent fba0642 commit 218ff20

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

tests/runai_model_streamer_test/test_weight_utils.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,11 @@ def test_runai_model_loader():
2323
runai_model_streamer_tensors = {}
2424
hf_safetensors_tensors = {}
2525

26-
for name, tensor in runai_safetensors_weights_iterator(safetensors):
26+
for name, tensor in runai_safetensors_weights_iterator(safetensors,
27+
True):
2728
runai_model_streamer_tensors[name] = tensor
2829

29-
for name, tensor in safetensors_weights_iterator(safetensors):
30+
for name, tensor in safetensors_weights_iterator(safetensors, True):
3031
hf_safetensors_tensors[name] = tensor
3132

3233
assert len(runai_model_streamer_tensors) == len(hf_safetensors_tensors)

0 commit comments

Comments
 (0)