Try installing triton only if linux & x86_64 (openai#1051)

zackees · May 5, 2023 · 74069fd · 74069fd
1 parent 1a5bb2a
commit 74069fd
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -1,4 +1,5 @@
 import os
+import platform
 import sys
 
 import pkg_resources
@@ -12,8 +13,8 @@
 
 
 requirements = []
-if sys.platform.startswith("linux"):
- triton_requirement = "triton>=2.0.0.dev20221202"
+if sys.platform.startswith("linux") and platform.machine() == "x86_64":
+ triton_requirement = "triton==2.0.0"
  try:
  import re
  import subprocess