Imported Upstream version 1.25.0

[platform/core/ml/nnfw.git] / compiler / dio-hdf5 / src / HDF5Importer.cpp

diff --git a/compiler/dio-hdf5/src/HDF5Importer.cpp b/compiler/dio-hdf5/src/HDF5Importer.cpp
index 9ae556b..9208990 100644 (file)
--- a/compiler/dio-hdf5/src/HDF5Importer.cpp
+++ b/compiler/dio-hdf5/src/HDF5Importer.cpp
@@ -128,16 +128,20 @@ int32_t HDF5Importer::numInputs(int32_t record_idx)
   return records.getNumObjs();
 }
 
-void HDF5Importer::readTensor(int32_t record_idx, int32_t input_idx, void *buffer)
+void HDF5Importer::readTensor(int32_t record_idx, int32_t input_idx, void *buffer,
+                              size_t buffer_bytes)
 {
   auto record = _group.openGroup(std::to_string(record_idx));
   auto tensor = record.openDataSet(std::to_string(input_idx));
 
+  if (tensor.getInMemDataSize() != buffer_bytes)
+    throw std::runtime_error("Buffer size does not match with the size of tensor data");
+
   readTensorData(tensor, static_cast<uint8_t *>(buffer));
 }
 
 void HDF5Importer::readTensor(int32_t record_idx, int32_t input_idx, DataType *dtype, Shape *shape,
-                              void *buffer)
+                              void *buffer, size_t buffer_bytes)
 {
   auto record = _group.openGroup(std::to_string(record_idx));
   auto tensor = record.openDataSet(std::to_string(input_idx));
@@ -148,6 +152,9 @@ void HDF5Importer::readTensor(int32_t record_idx, int32_t input_idx, DataType *d
   auto tensor_shape = tensor.getSpace();
   *shape = toInternalShape(tensor_shape);
 
+  if (tensor.getInMemDataSize() != buffer_bytes)
+    throw std::runtime_error("Buffer size does not match with the size of tensor data");
+
   switch (*dtype)
   {
     case DataType::FLOAT32: