Merge pull request #49541 from geetachavan1/cherrypicks_IXSSW

Fix FPE issue with `tf.raw_ops.FusedBatchNorm`.
tensorflow · May 27, 2021 · 99931bc · 99931bc
2 parents 636023e + 773ff26
commit 99931bc
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cc b/tensorflow/core/kernels/fused_batch_norm_op.cc
@@ -293,6 +293,9 @@ struct FusedBatchNorm<CPUDevice, T, U, /* is_training= */ false> {
     const CPUDevice& d = context->eigen_device<CPUDevice>();
 
     const int depth = x.dimension(3);
+    OP_REQUIRES(
+        context, depth != 0,
+        errors::Internal("The 4th element in the input shape cannot be 0."));
     const int size = x.size();
     const int rest_size = size / depth;
     Eigen::DSizes<Eigen::Index, 2> rest_by_depth(rest_size, depth);