Fixing issue #552

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

Author: borisfom

core/conversion/converters/impl/batch_norm.cpp

@@ -40,6 +40,7 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       LOG_DEBUG("momentum disregarded");
       LOG_DEBUG("training disregarded");
       LOG_DEBUG("cudnn disregarded");
+      TRTORCH_CHECK(orig_shape.nbDims > 2 , "Unable to create batch normalization layer from node: " << *n);
 
       // Expand spatial dims from 1D to 2D if needed
       bool expandDims = (orig_shape.nbDims < 4);

core/conversion/converters/impl/conv_deconv.cpp

@@ -30,7 +30,8 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
   LOG_DEBUG("out_padding: " << out_padding);
   LOG_DEBUG("groups: " << groups);
 
-  // Expand spatial dims from 1D to 2D if needed
+  TRTORCH_CHECK(orig_dims.nbDims > 2 , "Unable to create convolution layer from node: " << *n);
+
   bool expandDims = (orig_dims.nbDims < 4);
   if (expandDims) {
     in = addPadding(ctx, n, in, 4);

core/conversion/converters/impl/pooling.cpp

@@ -48,7 +48,7 @@ bool AdaptivePoolingConverter(
 
   auto orig_dims = in->getDimensions();
   bool expandDims = (orig_dims.nbDims < 4);
-
+  TRTORCH_CHECK(orig_dims.nbDims > 2 , "Unable to create pooling layer from node: " << *n);
   if (expandDims) {
     in = addPadding(ctx, n, in, 4, false, false);
   }
@@ -122,6 +122,7 @@ bool PoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args,
 
   // Max Pool needs at least 4D input
   auto orig_dims = in->getDimensions();
+  TRTORCH_CHECK(orig_dims.nbDims > 2 , "Unable to create pooling layer from node: " << *n);
   bool expandDims = (orig_dims.nbDims < 4);
 
   if (expandDims) {