Install fa2

pytorch-labs · Oct 17, 2024 · be243ae · be243ae
1 parent 8a735f3
commit be243ae
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/install.py b/install.py
@@ -115,11 +115,7 @@ def install_tk():
     if args.fbgemm or args.all:
         logger.info("[tritonbench] installing FBGEMM...")
         install_fbgemm()
-    # TODO: for some reason, fa2 compile will break docker build
-    if args.fa2:
-        logger.info("[tritonbench] installing fa2...")
-        install_fa2()
-    if args.fa2_compile:
+    if args.fa2 or args.all:
         logger.info("[tritonbench] installing fa2 from source...")
         install_fa2(compile=True)
     if args.fa3 or args.all: