fix device bug in async mempool
@@ -1600,33 +1600,32 @@ cdef class MemoryAsyncPool:
def __init__(self, pool_handles='current'):
_util.experimental('cupy.cuda.MemoryAsyncPool')
- cdef int dev_id, dev_counts
+ cdef int dev_id, prev_dev_id, dev_counts
cdef dict limit = _parse_limit_string()
dev_counts = runtime.getDeviceCount()
self._pools = []
self.memoryAsyncHasStat = (runtime.driverGetVersion() >= 11030)
+ prev_dev_id = runtime.getDevice()
if (cpython.PySequence_Check(pool_handles)
and not isinstance(pool_handles, str)):
# allow different kinds of handles on each device
for dev_id in range(dev_counts):
- prev_device = runtime.getDevice()
try:
runtime.setDevice(dev_id)
self._pools.append(self.set_pool(
pool_handles[dev_id], dev_id))
self.set_limit(**limit)
finally:
- runtime.setDevice(dev_id)
+ runtime.setDevice(prev_dev_id)
else:
# use the same argument for all devices
for dev_id in range(dev_counts):
- prev_device = runtime.getDevice()
try:
runtime.setDevice(dev_id)
self._pools.append(self.set_pool(pool_handles, dev_id))
self.set_limit(**limit)
finally:
- runtime.setDevice(dev_id)
+ runtime.setDevice(prev_dev_id)
cdef intptr_t set_pool(self, handle, int dev_id) except? 0:
cdef intptr_t pool
