Memory Leak Fixes in NNTrainer

This PR addresses memory leak issues in the NNTrainer. **Changes proposed in this PR:** - Release MemoryData pointer in the Tensor constructor using the vector. **Self-evaluation:** 1. Build test: [X]Passed [ ]Failed [ ]Skipped 2. Run test: [X]Passed [ ]Failed [ ]Skipped Signed-off-by: Donghyeon Jeong <[email protected]>
skykongkong8 · Feb 7, 2025 · a6ce493 · a6ce493
1 parent a497055
commit a6ce493
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 2 deletions.
diff --git a/nntrainer/tensor/char_tensor.cpp b/nntrainer/tensor/char_tensor.cpp
@@ -69,6 +69,7 @@ CharTensor::CharTensor(
     (void *)(new int8_t[dim.getDataLen() + sizeof(float) * scale_size()]()));
   data = std::shared_ptr<MemoryData>(mem_data, [](MemoryData *mem_data) {
     delete[] mem_data->getAddr<int8_t>();
+    delete mem_data;
   });
 
   offset = 0;

diff --git a/nntrainer/tensor/float_tensor.h b/nntrainer/tensor/float_tensor.h
@@ -88,6 +88,7 @@ class FloatTensor : public TensorBase {
       new MemoryData((void *)(new float[dim.getDataLen()]()));
     data = std::shared_ptr<MemoryData>(mem_data, [](MemoryData *mem_data) {
       delete[] mem_data->getAddr<float>();
+      delete mem_data;
     });
 
     offset = 0;

diff --git a/nntrainer/tensor/half_tensor.h b/nntrainer/tensor/half_tensor.h
@@ -87,6 +87,7 @@ class HalfTensor : public TensorBase {
       new MemoryData((void *)(new _FP16[dim.getDataLen()]()));
     data = std::shared_ptr<MemoryData>(mem_data, [](MemoryData *mem_data) {
       delete[] mem_data->getAddr<_FP16>();
+      delete mem_data;
     });
 
     offset = 0;

diff --git a/nntrainer/tensor/short_tensor.cpp b/nntrainer/tensor/short_tensor.cpp
@@ -67,6 +67,7 @@ ShortTensor::ShortTensor(
                          sizeof(float) / sizeof(int16_t) * scale_size()]()));
   data = std::shared_ptr<MemoryData>(mem_data, [](MemoryData *mem_data) {
     delete[] mem_data->getAddr<int16_t>();
+    delete mem_data;
   });
 
   offset = 0;

diff --git a/nntrainer/tensor/uint_tensor.cpp b/nntrainer/tensor/uint_tensor.cpp
@@ -61,8 +61,10 @@ UIntTensor<T>::UIntTensor(
   initializer = Initializer::NONE;
 
   MemoryData *mem_data = new MemoryData((void *)(new T[dim.getDataLen()]()));
-  data = std::shared_ptr<MemoryData>(
-    mem_data, [](MemoryData *mem_data) { delete[] mem_data->getAddr<T>(); });
+  data = std::shared_ptr<MemoryData>(mem_data, [](MemoryData *mem_data) {
+    delete[] mem_data->getAddr<T>();
+    delete mem_data;
+  });
 
   offset = 0;