fix: test regression

sehoffmann · Apr 2, 2024 · 574e904 · 574e904
1 parent 4276453
commit 574e904
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/test/test_data.py b/test/test_data.py
@@ -152,9 +152,10 @@ def test_shuffled(self):
         world_size = 3
         chunk_size = 15
 
-        chunks_1 = list(sharded_xr_dataset(ds, chunk_size, 'x', world_size=world_size, rank=0, shuffle=True, seed=0))
-        chunks_2 = list(sharded_xr_dataset(ds, chunk_size, 'x', world_size=world_size, rank=1, shuffle=True, seed=0))
-        chunks_3 = list(sharded_xr_dataset(ds, chunk_size, 'x', world_size=world_size, rank=2, shuffle=True, seed=0))
+        shard = partial(sharded_xr_dataset, ds, 'x', chunk_size, world_size=world_size, shuffle=True, seed=0)
+        chunks_1 = list(shard(rank=0))
+        chunks_2 = list(shard(rank=1))
+        chunks_3 = list(shard(rank=2))
 
         assert len(chunks_1) == 2
         assert len(chunks_2) == 2