From f14e4beeb5e421c74ca58822e41f01423611a49a Mon Sep 17 00:00:00 2001 From: raver119 Date: Fri, 24 Jan 2020 15:37:24 +0300 Subject: [PATCH] missing countOut for GPU Signed-off-by: raver119 --- libnd4j/include/array/cuda/DataBuffer.cu | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/libnd4j/include/array/cuda/DataBuffer.cu b/libnd4j/include/array/cuda/DataBuffer.cu index 7cb66ac9d..28e0c432f 100644 --- a/libnd4j/include/array/cuda/DataBuffer.cu +++ b/libnd4j/include/array/cuda/DataBuffer.cu @@ -132,6 +132,12 @@ void DataBuffer::deleteSpecial() { RELEASE_SPECIAL(p, _workspace); _specialBuffer = nullptr; _isOwnerSpecial = false; + + // count out towards DataBuffer device, only if we're not in workspace + if (_workspace == nullptr) { + nd4j::memory::MemoryCounter::getInstance()->countOut(_deviceId, getLenInBytes()); + nd4j::memory::MemoryCounter::getInstance()->countOut(nd4j::memory::MemoryType::DEVICE, getLenInBytes()); + } } }