26 lines
1.1 KiB
Diff
26 lines
1.1 KiB
Diff
|
From 1a2a87229d1d61e23a39373777c056161eb4084d Mon Sep 17 00:00:00 2001
|
||
|
|
From: Amit Patankar <amitpatankar@google.com>
|
||
|
|
Date: Wed, 28 Apr 2021 11:30:28 -0700
|
||
|
|
Subject: [PATCH] Fix FPE issue with `tf.raw_ops.FusedBatchNorm`.
|
||
|
|
|
||
|
|
PiperOrigin-RevId: 370948185
|
||
|
|
Change-Id: If0c8e0320062ed6363e94ff5fe38e6a301f69ac2
|
||
|
|
---
|
||
|
|
tensorflow/core/kernels/fused_batch_norm_op.cc | 3 +++
|
||
|
|
1 file changed, 3 insertions(+)
|
||
|
|
|
||
|
|
diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cc b/tensorflow/core/kernels/fused_batch_norm_op.cc
|
||
|
|
index 9b917348e2afa..e564b19857c38 100644
|
||
|
|
--- 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);
|