aitemplate/AITemplate (#978)

Summary: Pull Request resolved: #978 Reviewed By: zsol Differential Revision: D52193552 fbshipit-source-id: e704d80114ff21dcf52fa7dd7962bad92144377f
facebookincubator · Dec 15, 2023 · da97807 · da97807
1 parent b29432a
commit da97807
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/aitemplate/utils/mk_cutlass_lib/mk_cutlass_lib.py b/python/aitemplate/utils/mk_cutlass_lib/mk_cutlass_lib.py
@@ -87,11 +87,15 @@ 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")
 
     assert output_path is not None
     assert cutlass_path is not None
 
     mk_cutlass_lib(cutlass_path + "/cutlass", os.path.dirname(output_path))
+
+
+if __name__ == "__main__":
+    main()  # pragma: no cover