Merge pull request #49838 from tensorflow/cp_2_3

[CherryPick]Fix 2 issues with .
tensorflow · May 31, 2021 · 4b31fcf · 4b31fcf
2 parents 28638f5 + 7a46fb2
commit 4b31fcf
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/tensorflow/core/kernels/conv_ops_3d.h b/tensorflow/core/kernels/conv_ops_3d.h
@@ -56,6 +56,11 @@ struct LaunchConvOp<CPUDevice, T, OpKernelContextT> {
                 errors::InvalidArgument("CPU implementation of Conv3D "
                                         "currently only supports dilated rates "
                                         "of 1."));
+    OP_REQUIRES(context, filter.dim_size(3) == input.dim_size(input.dims() - 1),
+                errors::InvalidArgument(
+                    "Number of channels in filter (", filter.dim_size(3),
+                    ") must match last dimension of input (",
+                    input.dim_size(input.dims() - 1), ")"));
     functor::CuboidConvolution<CPUDevice, T>()(
         context->template eigen_device<CPUDevice>(), output->tensor<T, 5>(),
         input.tensor<T, 5>(), filter.tensor<T, 5>(), strides[2], strides[1],
@@ -135,6 +140,8 @@ class Conv3DOp : public BinaryOpBase<T, OpKernelT, OpKernelConstructionT> {
     const int64 filter_depth = filter.dim_size(3);
     const int64 out_depth = filter.dim_size(4);
 
+    OP_REQUIRES(context, filter_depth != 0,
+                errors::InvalidArgument("filter_depth must be non-zero"));
     OP_REQUIRES(context, in_depth % filter_depth == 0,
                 errors::InvalidArgument(
                     "Input depth must be evenly divisible by filter depth: ",