Call gc to release memory when memory is not enough
import collections
import ctypes
+import gc
import warnings
import weakref
@@ -358,6 +359,12 @@ cdef class SingleDeviceMemoryPool:
if e.status != runtime.errorMemoryAllocation:
raise
self.free_all_free()
+ try:
+ mem = self._alloc(size).mem
+ except runtime.CUDARuntimeError as e:
+ if e.status != runtime.errorMemoryAllocation:
+ raise
+ gc.collect()
mem = self._alloc(size).mem
self._in_use[mem.ptr] = mem
