missing countOut for GPU

Signed-off-by: raver119 <raver119@gmail.com>
master
raver119 2020-01-24 15:37:24 +03:00
parent 99cfb88933
commit f14e4beeb5
1 changed files with 6 additions and 0 deletions

View File

@ -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());
}
}
}