Skip to content

Commit d79a41b

Browse files
xwhzzLeiWang1999
andauthored
[BugFix] Correct argparse for example_convolution test (#474)
* add convolution example to CI * lint fix * Update test_example_convolution.py * fix bug --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
1 parent 4ce2b7a commit d79a41b

File tree

2 files changed

+3
-3
lines changed

2 files changed

+3
-3
lines changed

examples/convolution/example_convolution.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -232,7 +232,7 @@ def main(
232232
return main
233233

234234

235-
def main():
235+
def main(argv=None):
236236
parser = argparse.ArgumentParser()
237237
parser.add_argument('--n', type=int, default=128, help='n')
238238
parser.add_argument('--c', type=int, default=128, help='c')
@@ -254,7 +254,7 @@ def main():
254254
default=True,
255255
help="Whether to enable BitBLAS roller for search space")
256256

257-
args = parser.parse_args()
257+
args = parser.parse_args(argv)
258258
N, C, H, W, F, K, S, D, P = args.n, args.c, args.h, args.w, args.f, args.k, args.s, args.d, args.p
259259
a = torch.randn(N, H, W, C).cuda().half()
260260
b = torch.randn(K, K, C, F).cuda().half()

examples/convolution/test_example_convolution.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
@tilelang.testing.requires_cuda
99
def test_example_convolution():
10-
example_convolution.main()
10+
example_convolution.main([])
1111

1212

1313
if __name__ == "__main__":

0 commit comments

Comments
 (0)