fix test

Eventual-Inc · Oct 5, 2024 · 6db6229 · 6db6229
1 parent 127807e
commit 6db6229
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/expressions/test_stateful_udf.py b/tests/expressions/test_stateful_udf.py
@@ -142,6 +142,8 @@ def __call__(self, data):
 
             return [self.cuda_visible_devices] * len(data)
 
+    FractionalGpuUdf = FractionalGpuUdf.with_concurrency(2)
+
     df = daft.from_pydict({"x": [1, 2]})
     df = df.into_partitions(2)
     df = df.select(FractionalGpuUdf(df["x"]))