[ahub] fix ahub issue

[platform/core/ml/nntrainer.git] / nntrainer / tensor / tensor.cpp

diff --git a/nntrainer/tensor/tensor.cpp b/nntrainer/tensor/tensor.cpp
index 63f7e48..f76b13f 100644 (file)
--- a/nntrainer/tensor/tensor.cpp
+++ b/nntrainer/tensor/tensor.cpp
@@ -1600,8 +1600,12 @@ void Tensor::save(std::ostream &file) {
   NNTR_THROW_IF(!contiguous, std::invalid_argument)
     << getName() << " is not contiguous, cannot save.";
 
-  checkedWrite(file, (char *)getData(), bytes(),
-               "[Tensor::save] operation failed");
+  std::streamsize sz = static_cast<std::streamsize>(bytes());
+  NNTR_THROW_IF(sz < 0, std::invalid_argument)
+    << "save size: " << bytes()
+    << " is too big. It cannot be represented by std::streamsize";
+
+  checkedWrite(file, (char *)getData(), sz, "[Tensor::save] operation failed");
   putData();
 }
 
@@ -1609,8 +1613,13 @@ void Tensor::read(std::ifstream &file) {
   NNTR_THROW_IF(!contiguous, std::invalid_argument)
     << getName() << " is not contiguous, cannot read.";
 
-  checkedRead(file, (char *)getData(), bytes(),
-              "[Tensor::read] operation failed");
+  std::streamsize sz = static_cast<std::streamsize>(bytes());
+
+  NNTR_THROW_IF(sz < 0, std::invalid_argument)
+    << "read size: " << bytes()
+    << " is too big. It cannot be represented by std::streamsize";
+
+  checkedRead(file, (char *)getData(), sz, "[Tensor::read] operation failed");
   putData();
 }