Skip to content

Commit 8c680b5

Browse files
minor: change device selection in poststack_cupy.py
Co-authored-by: Rohan Babbar <[email protected]>
1 parent 75d7b10 commit 8c680b5

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tutorials_cupy/poststack_cupy.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@
2020

2121
plt.close("all")
2222
rank = MPI.COMM_WORLD.Get_rank()
23-
cp.cuda.Device(device=rank).use();
23+
device_count = cp.cuda.runtime.getDeviceCount()
24+
cp.cuda.Device(rank % device_count).use()
2425

2526
###############################################################################
2627
# Let's start by defining all the parameters required by the

0 commit comments

Comments
 (0)