Skip to content
Permalink
Browse files Browse the repository at this point in the history
Ensure num_streams >= 0 in tf.raw_ops.BoostedTreesCreateQuantileStrea…
…mResource

PiperOrigin-RevId: 387452765
Change-Id: I9990c760e177fabca6a3b9b4612ceeaeeba51495
  • Loading branch information
pak-laura authored and tensorflower-gardener committed Jul 28, 2021
1 parent f8a1ac8 commit 8a84f7a
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/kernels/boosted_trees/quantile_ops.cc
Expand Up @@ -116,6 +116,9 @@ class BoostedTreesCreateQuantileStreamResourceOp : public OpKernel {
const Tensor* num_streams_t;
OP_REQUIRES_OK(context, context->input(kNumStreamsName, &num_streams_t));
int64_t num_streams = num_streams_t->scalar<int64>()();
OP_REQUIRES(context, num_streams >= 0,
errors::InvalidArgument(
"Num_streams input cannot be a negative integer"));

auto result =
new QuantileStreamResource(epsilon, max_elements_, num_streams);
Expand Down

0 comments on commit 8a84f7a

Please sign in to comment.