Skip to content
Permalink
Browse files Browse the repository at this point in the history
Fix 2 issues with Conv3D.
We have an issue where the dimensions are not matching and this causes Eigen to crash on an assert.

Then, we have an issue where we accidentally do a division by 0.

PiperOrigin-RevId: 369242785
Change-Id: Ie94067b2d41f58699af99ebb5af335ad9defd931
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Apr 19, 2021
1 parent 4203360 commit 799f835
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions tensorflow/core/kernels/conv_ops_3d.cc
Expand Up @@ -69,6 +69,11 @@ struct LaunchConvOp<CPUDevice, T> {
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->eigen_device<CPUDevice>(), output->tensor<T, 5>(),
input.tensor<T, 5>(), filter.tensor<T, 5>(), strides[2], strides[1],
Expand Down Expand Up @@ -142,6 +147,8 @@ class Conv3DOp : public BinaryOp<T> {
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: ",
Expand Down

0 comments on commit 799f835

Please sign in to comment.