Add atexit code
@@ -12,6 +12,7 @@ from cupy.cuda import cudnn
_cudnn_version = cudnn.getVersion()
_handles = {}
+_nd_tensor_cache = {}
def get_handle():
@@ -33,6 +34,11 @@ def reset_handles():
cudnn.destroy(handle)
+@atexit.register
+def _clear_cache():
+ _nd_tensor_cache.clear()
+
+
class Descriptor(object):
def __init__(self, descriptor, destroyer):
@@ -88,9 +94,6 @@ def create_uninitialized_tensor_descriptor():
return desc
-_nd_tensor_cache = {}
-
-
def create_tensor_nd_descriptor(arr):
desc = Descriptor(cudnn.createTensorDescriptor(),
cudnn.destroyTensorDescriptor)
