Merge pull request #49653 from geetachavan1/cherrypicks_UE6M7

Add missing validations to reverse_sequence_op
tensorflow · May 26, 2021 · 0f90ce8 · 0f90ce8
2 parents bf796c4 + c42d04b
commit 0f90ce8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/reverse_sequence_op.cc b/tensorflow/core/kernels/reverse_sequence_op.cc
@@ -115,6 +115,10 @@ class ReverseSequenceOp : public OpKernel {
       : OpKernel(context) {
     OP_REQUIRES_OK(context, context->GetAttr("batch_dim", &batch_dim_));
     OP_REQUIRES_OK(context, context->GetAttr("seq_dim", &seq_dim_));
+    OP_REQUIRES(context, batch_dim_ >= 0,
+                errors::InvalidArgument("Invalid batch_dim ", batch_dim_));
+    OP_REQUIRES(context, seq_dim_ >= 0,
+                errors::InvalidArgument("Invalid seq_dim ", seq_dim_));
   }
 
   void Compute(OpKernelContext* context) override {