diff --git a/python/aitemplate/utils/mk_cutlass_lib/mk_cutlass_lib.py b/python/aitemplate/utils/mk_cutlass_lib/mk_cutlass_lib.py
index 8231bc7b6..5b39b0809 100644
--- a/python/aitemplate/utils/mk_cutlass_lib/mk_cutlass_lib.py
+++ b/python/aitemplate/utils/mk_cutlass_lib/mk_cutlass_lib.py
@@ -87,7 +87,7 @@ def process_code(src_path, dst_path, code_set):
     return dst_prefix
 
 
-if __name__ == "__main__":
+def main() -> None:
     cutlass_path = os.getenv("SRCDIR")
     output_path = os.getenv("OUT")
 
@@ -95,3 +95,7 @@ def process_code(src_path, dst_path, code_set):
     assert cutlass_path is not None
 
     mk_cutlass_lib(cutlass_path + "/cutlass", os.path.dirname(output_path))
+
+
+if __name__ == "__main__":
+    main()  # pragma: no cover