From 79865b542f9ffdc9caeb255631f7c56f1d4b6517 Mon Sep 17 00:00:00 2001 From: Amit Patankar Date: Tue, 4 May 2021 18:05:46 -0700 Subject: [PATCH] Fix memory corruption issue with `tf.raw_ops.DrawBoundingBoxesV2`. --- tensorflow/core/kernels/draw_bounding_box_op.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorflow/core/kernels/draw_bounding_box_op.cc b/tensorflow/core/kernels/draw_bounding_box_op.cc index 30de99b7..1e866325 100644 --- a/tensorflow/core/kernels/draw_bounding_box_op.cc +++ b/tensorflow/core/kernels/draw_bounding_box_op.cc @@ -73,6 +73,12 @@ class DrawBoundingBoxesOp : public OpKernel { errors::InvalidArgument("Channel depth should be either 1 (GRY), " "3 (RGB), or 4 (RGBA)")); + OP_REQUIRES( + context, boxes.dim_size(2) == 4, + errors::InvalidArgument( + "The size of the third dimension of the box mustbe 4. Received: ", + boxes.dim_size(2))); + const int64 batch_size = images.dim_size(0); const int64 height = images.dim_size(1); const int64 width = images.dim_size(2); -- 2.23.0