Pin triton version to avoid import error (#642)

Transformer engine seems to be having an issue with triton 3.2.0, pinning to 3.1.0 until we update to a new NGC base image Signed-off-by: Peter St. John <[email protected]>
NVIDIA · Jan 23, 2025 · 17c6b20 · 17c6b20
1 parent 49a72a9
commit 17c6b20
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/requirements-test.txt b/requirements-test.txt
@@ -13,3 +13,6 @@ pyfaidx==0.8.1.3
 # See https://nvidia.slack.com/archives/C02A7LYGHK8/p1734727482697309
 pytorch-lightning<2.5.0
 lightning<2.5.0
+
+# Temporary pin for triton
+triton<=3.1.0