Commit 3d7da362 by zacario-li Committed by Tianqi Chen

fix topi c++ conv2d_nchw lambda expr issue (#3570)

parent 2973f8a6
...@@ -268,7 +268,7 @@ inline tvm::Tensor conv2d_nchw(const tvm::Tensor& I, ...@@ -268,7 +268,7 @@ inline tvm::Tensor conv2d_nchw(const tvm::Tensor& I,
: pad(I, {tvm::Expr(0), tvm::Expr(0), pad_h, pad_w}); : pad(I, {tvm::Expr(0), tvm::Expr(0), pad_h, pad_w});
auto l = [&](tvm::Var b, tvm::Var o, tvm::Var h, tvm::Var w) { auto l = [&](tvm::Var b, tvm::Var o, tvm::Var h, tvm::Var w) {
return tvm::sum( return tvm::sum(
T(b, i, stride_h * h + kh, stride_w * w + kw) * W(i, o, kh, kw), T(b, i, stride_h * h + kh, stride_w * w + kw) * W(o, i, kh, kw),
{i, kh, kw}); {i, kh, kw});
}; };
return tvm::compute(output_shape, l, name, tag); return tvm::compute(output_shape, l, name, tag);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment