pin torch version in amplify

pstjohn · pstjohn · commit b35df98b8c87 · 2025-09-04T10:08:29.000-07:00
Signed-off-by: Peter St. John &lt;pstjohn@nvidia.com&gt;
diff --git a/models/amplify/pyproject.toml b/models/amplify/pyproject.toml
@@ -19,8 +19,8 @@ dependencies = [
     "nvidia_resiliency_ext",
     "omegaconf",
     "pytest",
-    "torch",
-    # "transformer_engine[pytorch]",
+    "torch==2.6.0a0+ecf3bae40a.nv25.01",
+    "transformer_engine[pytorch]",
     "transformers",
     "xformers",
 ]
