Merge pull request #49643 from geetachavan1/cherrypicks_SPLNT

Fix memory corruption issue with `tf.raw_ops.DrawBoundingBoxesV2`.
tensorflow · May 26, 2021 · 70d5066 · 70d5066
2 parents db824f3 + 29b4096
commit 70d5066
Showing 1 changed file with 6 additions and 0 deletions.
diff --git 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 must be 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);