Skip to content

Commit 7c3d68c

Browse files
committed
TST: fix cupy to_device test on multiple devices
1 parent 2b5e289 commit 7c3d68c

File tree

1 file changed

+8
-6
lines changed

1 file changed

+8
-6
lines changed

tests/test_cupy.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,17 @@
88
def test_to_device_with_stream():
99
devices = xp.__array_namespace_info__().devices()
1010
streams = [
11-
Stream(),
12-
Stream(non_blocking=True),
13-
Stream(null=True),
14-
Stream(ptds=True),
15-
123, # dlpack stream
11+
lambda: Stream(),
12+
lambda: Stream(non_blocking=True),
13+
lambda: Stream(null=True),
14+
lambda: Stream(ptds=True),
15+
lambda: 123, # dlpack stream
1616
]
1717

1818
a = xp.asarray([1, 2, 3])
1919
for dev in devices:
20-
for stream in streams:
20+
for stream_gen in streams:
21+
with dev:
22+
stream = stream_gen()
2123
b = to_device(a, dev, stream=stream)
2224
assert device(b) == dev

0 commit comments

Comments
 (0)