From c03ad1a46d5b3f23df67dad03185a0ee16020c96 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Tue, 11 May 2021 18:36:43 -0700
Subject: [PATCH] Ensure validation sticks in banded_triangular_solve_op

PiperOrigin-RevId: 373275480
Change-Id: Id7717cf275b2d6fdb9441fbbe166d555182d2e79
---
 tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc b/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc
index c079c63b778ac7..b719f55b507b02 100644
--- a/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc
+++ b/tensorflow/core/kernels/linalg/banded_triangular_solve_op.cc
@@ -217,6 +217,7 @@ class BandedTriangularSolveOpCpu : public OpKernel {
     const Tensor& in1 = ctx->input(1);
 
     ValidateInputTensors(ctx, in0, in1);
+    if (!ctx->status().ok()) return;
 
     MatMulBCast bcast(in0.shape().dim_sizes(), in1.shape().dim_sizes());
     OP_REQUIRES(