Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent heap oob access in resource_variable_ops.cc
PiperOrigin-RevId: 387936433
Change-Id: I9e71ddaa8dbd51ec6afbf163a6b3b591f193b4f6
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 31, 2021
1 parent d5f28c9 commit bc9c546
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions tensorflow/core/kernels/resource_variable_ops.cc
Expand Up @@ -660,6 +660,11 @@ class ResourceGatherOp : public OpKernel {
OP_REQUIRES(
c, TensorShapeUtils::IsVectorOrHigher(params.shape()),
errors::InvalidArgument("params must be at least 1 dimensional"));
OP_REQUIRES(
c, params.shape().dims() >= batch_dims_,
errors::InvalidArgument("params must have at least ", batch_dims_,
" (batch_dims) dimensions but it has shape ",
params.shape().DebugString()));

// Check that we have enough index space
const int64_t N = indices.NumElements();
Expand Down

0 comments on commit bc9c546

Please sign in to comment.