From 4d1d08897a94762cd4b0f74770b0bc462ca3bd5d Mon Sep 17 00:00:00 2001 From: Lei Wen Date: Thu, 19 Oct 2023 16:17:58 +0800 Subject: [PATCH] bad code show case Try export onnx file Signed-off-by: John --- python/tvm/contrib/cutlass/build.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index cd990d0eea..283e903405 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -47,6 +47,20 @@ def _get_cutlass_path(): raise AssertionError(f"The CUTLASS root directory not found in: {invalid_paths}") +def ddd(): + print + + +class SubTreeExtractor(nn.Module): + def __init__(self): + super(SubTreeExtractor, self).__init__() + + def forward(self, x, y): + return x + y + +torch.onnx.export(SubTreeExtractor(), (torch.Tensor(1,2,3,4), 1), "1.onnx") + + def _get_cutlass_compile_options(sm, threads, use_fast_math=False): cutlass_root = _get_cutlass_path() cutlass_include = os.path.join(cutlass_root, "include")