Use stream to allocate cupy memory in new stream test

NVIDIA · Jan 17, 2025 · 06df21e · 06df21e
1 parent 6a1f2a8
commit 06df21e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/cuda_parallel/tests/test_reduce.py b/python/cuda_parallel/tests/test_reduce.py
@@ -584,7 +584,8 @@ def add_op(x, y):
         h_init=h_init,
         stream=stream_wrapper,
     )
-    d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
+    with stream:
+        d_temp_storage = cp.empty(temp_storage_size, dtype=np.uint8)
 
     reduce_into(d_temp_storage, d_in, d_out, d_in.size, h_init, stream=stream_wrapper)
     with stream: