Commit dc4320af authored by Fenglei's avatar Fenglei Committed by Robert Kimball

using create_gpu_buffer and free_gpu_buffer (#1930)

parent 96875313
...@@ -44,7 +44,7 @@ runtime::gpu::GPUTensor::GPUTensor(const ngraph::element::Type& element_type, ...@@ -44,7 +44,7 @@ runtime::gpu::GPUTensor::GPUTensor(const ngraph::element::Type& element_type,
} }
else if (m_buffer_size > 0) else if (m_buffer_size > 0)
{ {
CUDA_RT_SAFE_CALL(cudaMalloc(static_cast<void**>(&m_allocated_buffer_pool), m_buffer_size)); m_allocated_buffer_pool = runtime::gpu::create_gpu_buffer(m_buffer_size);
} }
} }
...@@ -57,7 +57,7 @@ runtime::gpu::GPUTensor::~GPUTensor() ...@@ -57,7 +57,7 @@ runtime::gpu::GPUTensor::~GPUTensor()
{ {
if (!m_custom_memory && (m_allocated_buffer_pool != nullptr)) if (!m_custom_memory && (m_allocated_buffer_pool != nullptr))
{ {
CUDA_RT_SAFE_CALL_NO_THROW(cudaFree(m_allocated_buffer_pool)); runtime::gpu::free_gpu_buffer(m_allocated_buffer_pool);
} }
} }
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment