Merge pull request #3329 from okuta/use-cupy-allocator

Use CuPy default memory pool
chainer · Nov 6, 2017 · fa55e1c · fa55e1c
2 parents 84c854e + bb0b47d
commit fa55e1c
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/chainer/cuda.py b/chainer/cuda.py
@@ -130,10 +130,9 @@ def __ne__(self, other):
 # Global states
 # ------------------------------------------------------------------------------
 if available:
-    memory_pool = cuda.MemoryPool()
-    cuda.set_allocator(memory_pool.malloc)
-    pinned_memory_pool = cuda.PinnedMemoryPool()
-    cuda.set_pinned_memory_allocator(pinned_memory_pool.malloc)
+    # This is for backward compatibility
+    memory_pool = cupy.get_default_memory_pool()
+    pinned_memory_pool = cupy.get_default_pinned_memory_pool()
 
 
 if six.PY2: