Skip to content
Permalink
Browse files Browse the repository at this point in the history
Validate num_elements input in tf.raw_ops.TensorListReserve
PiperOrigin-RevId: 383954564
Change-Id: I454bd78eff85bc4f16ddb7e608596971cca47f8f
  • Loading branch information
pak-laura authored and tensorflower-gardener committed Jul 10, 2021
1 parent 3e23241 commit 8a6e874
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions tensorflow/core/kernels/list_kernels.cc
Expand Up @@ -302,6 +302,10 @@ class TensorListReserve : public OpKernel {
PartialTensorShape element_shape;
OP_REQUIRES_OK(c, TensorShapeFromTensor(c->input(0), &element_shape));
int32 num_elements = c->input(1).scalar<int32>()();
OP_REQUIRES(c, num_elements >= 0,
errors::InvalidArgument("The num_elements to reserve must be a "
"non negative number, but got ",
num_elements));
TensorList output;
output.element_shape = element_shape;
output.element_dtype = element_dtype_;
Expand Down

0 comments on commit 8a6e874

Please sign in to comment.