[Minor] Fix test_cache.py CI test failure (#2684)

89efcf1c · Philipp Moritz · GitHub · c664b0e6 · 89efcf1c
Unverified Commit 89efcf1c authored Jan 31, 2024 by Philipp Moritz Committed by GitHub Jan 31, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

tests/kernels/test_cache.py tests/kernels/test_cache.py +2 -2

No files found.
--- a/tests/kernels/test_cache.py
+++ b/tests/kernels/test_cache.py
@@ -200,12 +200,12 @@ def test_swap_blocks(

    # Create the KV caches on the first device.
    src_key_caches, src_value_caches = kv_cache_factory(
-        num_blocks, block_size, 1, num_heads, head_size, dtype, seed,
+        num_blocks, block_size, 1, num_heads, head_size, dtype, None, seed,
        src_device)

    # Create the KV caches on the second device.
    dist_key_caches, dist_value_caches = kv_cache_factory(
-        num_blocks, block_size, 1, num_heads, head_size, dtype, seed,
+        num_blocks, block_size, 1, num_heads, head_size, dtype, None, seed,
        dst_device)

    src_key_caches_clone = src_key_caches[0].clone()