make sure cutn handle is cached properly
@@ -88,8 +88,10 @@ def _try_use_cutensornet(*args, **kwargs):
# prepare cutn inputs
device = cupy.cuda.runtime.getDevice()
- handle = cutn_handle_cache.get(
- device, cutensornet.create())
+ handle = cutn_handle_cache.get(device)
+ if handle is None:
+ handle = cutensornet.create()
+ cutn_handle_cache[device] = handle
cutn_options = {'device_id': device, 'handle': handle,
'memory_limit': 4**31} # TODO(leofang): fix?
