Extract channels from weight shape for conv2d. (#6805)

apache · Oct 31, 2020 · ceef616 · ceef616
1 parent 50fc938
commit ceef616
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/src/runtime/contrib/tensorrt/tensorrt_ops.cc b/src/runtime/contrib/tensorrt/tensorrt_ops.cc
@@ -242,7 +242,7 @@ class Conv2DOpConverter : public TensorRTOpConverter {
     auto str_dilation = params->node.GetAttr<std::vector<std::string>>("dilation");
     auto str_padding = params->node.GetAttr<std::vector<std::string>>("padding");
     int groups = std::stoi(params->node.GetAttr<std::vector<std::string>>("groups")[0]);
-    int channels = std::stoi(params->node.GetAttr<std::vector<std::string>>("channels")[0]);
+    int channels = weight_shape[0];
     // TRT conv2d op doesn't support asymmetric padding before 5.1, so we
     // workaround by adding a padding layer before the pooling op.
     nvinfer1::DimsHW prepadding, postpadding;

diff --git a/tests/python/contrib/test_tensorrt.py b/tests/python/contrib/test_tensorrt.py
@@ -251,7 +251,6 @@ def get_graph(
         out = relay.nn.conv2d(
             x,
             kernel,
-            channels=k_shape[0],
             kernel_size=k_shape[2:4],
             groups=groups,
             padding=padding,