42 lines
1.6 KiB
Diff
42 lines
1.6 KiB
Diff
From 106d8f4fb89335a2c52d7c895b7a7485465ca8d9 Mon Sep 17 00:00:00 2001
|
|
From: Mihai Maruseac <mihaimaruseac@google.com>
|
|
Date: Tue, 27 Apr 2021 17:47:36 -0700
|
|
Subject: [PATCH] Prevent division by 0 in TFLite
|
|
|
|
---
|
|
tensorflow/lite/kernels/depth_to_space.cc | 1 +
|
|
tensorflow/lite/kernels/depth_to_space_test.cc | 5 +++++
|
|
2 files changed, 6 insertions(+)
|
|
|
|
diff --git a/tensorflow/lite/kernels/depth_to_space.cc b/tensorflow/lite/kernels/depth_to_space.cc
|
|
index 1637ad43..c2047f10 100644
|
|
--- a/tensorflow/lite/kernels/depth_to_space.cc
|
|
+++ b/tensorflow/lite/kernels/depth_to_space.cc
|
|
@@ -58,6 +58,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
|
TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
|
|
|
|
const int block_size = params->block_size;
|
|
+ TF_LITE_ENSURE(context, block_size > 0);
|
|
const int input_height = input->dims->data[1];
|
|
const int input_width = input->dims->data[2];
|
|
const int input_channels = input->dims->data[3];
|
|
diff --git a/tensorflow/lite/kernels/depth_to_space_test.cc b/tensorflow/lite/kernels/depth_to_space_test.cc
|
|
index 4429faf9..e0de01a9 100644
|
|
--- a/tensorflow/lite/kernels/depth_to_space_test.cc
|
|
+++ b/tensorflow/lite/kernels/depth_to_space_test.cc
|
|
@@ -60,6 +60,11 @@ TEST(DepthToSpaceOpModel, BadBlockSize) {
|
|
EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4),
|
|
"Cannot allocate tensors");
|
|
}
|
|
+
|
|
+TEST(DepthToSpaceOpModel, NoBlockSize) {
|
|
+ EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0),
|
|
+ "Cannot allocate tensors");
|
|
+}
|
|
#endif
|
|
|
|
TEST(DepthToSpaceOpModel, Float32) {
|
|
--
|
|
2.23.0
|
|
|