fix output_shape in conv2d_nchw (#1613)

apache · Aug 17, 2018 · acc2151 · acc2151
1 parent e282915
commit acc2151
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/topi/include/topi/nn.h b/topi/include/topi/nn.h
@@ -265,7 +265,7 @@ inline tvm::Tensor conv2d_nchw(const tvm::Tensor& I,
   auto pW = I->shape[3];
   tvm::Array<tvm::Expr> output_shape{
       I->shape[0],                                            // B
-      W->shape[1],                                            // O
+      W->shape[0],                                            // O
       (I->shape[2] - W->shape[2] + 2 * pad_h) / stride_h + 1,  // H
       (I->shape[3] - W->shape[3] + 2 * pad_w) / stride_w + 1   // W
   };