Merge pull request #49591 from geetachavan1/cherrypicks_PUF1Y

Prevent check fail in FFT
tensorflow · May 25, 2021 · 711324e · 711324e
2 parents 19f3f1b + f116a60
commit 711324e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/fft_ops.cc b/tensorflow/core/kernels/fft_ops.cc
@@ -221,6 +221,9 @@ class FFTCPU : public FFTBase {
       input_slice_sizes[i] = fft_shape[i - 1];
       temp_shape.AddDim(fft_shape[i - 1]);
     }
+    OP_REQUIRES(ctx, temp_shape.num_elements() > 0,
+                errors::InvalidArgument("Obtained a FFT shape of 0 elements: ",
+                                        temp_shape.DebugString()));
 
     auto output = out->flat_inner_dims<ComplexT, FFTRank + 1>();
     const Eigen::DSizes<Eigen::DenseIndex, FFTRank + 1> zero_start_indices;