[bugfix] Resolve segfault in tensor apply

This PR fixes a bug where a segmentation fault occurs when the output tensor is empty. The fix initializes the output tensor when empty to avoid this error. **Self-evaluation:** 1. Build test: [X]Passed [ ]Failed [ ]Skipped 2. Run test: [X]Passed [ ]Failed [ ]Skipped Signed-off-by: Donghyeon Jeong <[email protected]>
nnstreamer · Jan 23, 2024 · 270b106 · 270b106
1 parent 07db2d5
commit 270b106
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/nntrainer/tensor/float_tensor.cpp b/nntrainer/tensor/float_tensor.cpp
@@ -261,6 +261,8 @@ void FloatTensor::initialize(Initializer init) {
 
 TensorV2 &FloatTensor::apply(std::function<float(float)> f,
                              TensorV2 &output) const {
+  CREATE_V2_IF_EMPTY_DIMS(output, dim, nullptr);
+
   if (contiguous && output.getContiguous()) {
     const float *data = (float *)getData();
     float *rdata = output.getData<float>();

diff --git a/nntrainer/tensor/half_tensor.cpp b/nntrainer/tensor/half_tensor.cpp
@@ -261,6 +261,8 @@ void HalfTensor::initialize(Initializer init) {
 
 TensorV2 &HalfTensor::apply(std::function<_FP16(_FP16)> f,
                             TensorV2 &output) const {
+  CREATE_V2_IF_EMPTY_DIMS(output, dim, nullptr);
+
   if (contiguous && output.getContiguous()) {
     const _FP16 *data = (_FP16 *)getData();
     _FP16 *rdata = output.getData<_FP16>();