diff options
Diffstat (limited to 'tensorflow/core/common_runtime/gpu/pool_allocator_test.cc')
| -rw-r--r-- | tensorflow/core/common_runtime/gpu/pool_allocator_test.cc | 6 |
1 files changed, 4 insertions, 2 deletions
diff --git a/tensorflow/core/common_runtime/gpu/pool_allocator_test.cc b/tensorflow/core/common_runtime/gpu/pool_allocator_test.cc index a4c8d5fe86..583bff2c07 100644 --- a/tensorflow/core/common_runtime/gpu/pool_allocator_test.cc +++ b/tensorflow/core/common_runtime/gpu/pool_allocator_test.cc @@ -15,8 +15,9 @@ limitations under the License. #if GOOGLE_CUDA -#include "tensorflow/core/common_runtime/gpu/pool_allocator.h" +#include "tensorflow/core/common_runtime/pool_allocator.h" +#include "tensorflow/core/common_runtime/gpu/cuda_host_allocator.h" #include "tensorflow/core/platform/stream_executor.h" #include "tensorflow/core/platform/test.h" @@ -96,7 +97,8 @@ TEST(PoolAllocatorTest, Alignment) { TEST(PoolAllocatorTest, AutoResize) { PoolAllocator pool(2 /*pool_size_limit*/, true /*auto_resize*/, - new BasicCPUAllocator, new NoopRounder, "pool"); + new BasicCPUAllocator(0 /*numa_node*/), new NoopRounder, + "pool"); // Alloc/dealloc 10 sizes just a few times, confirming pool size // stays at 2. |