Rocm (#293)

* ROCm changes * Update ping * - * Cleanup the rocm script * use rocm branch * - * Update run_rocm.sh --------- Co-authored-by: Your Name <[email protected]>
mila-iqia · Nov 21, 2024 · 0b93d59 · 0b93d59
1 parent 798a186
commit 0b93d59
Show file tree

Hide file tree

Showing 29 changed files with 2,723 additions and 696 deletions.
diff --git a/.pin/constraints-rocm-torch.txt b/.pin/constraints-rocm-torch.txt
diff --git a/benchmarks/brax/benchfile.py b/benchmarks/brax/benchfile.py
@@ -5,5 +5,9 @@ class BraxBenchmark(Package):
     base_requirements = "requirements.in"
     main_script = "main.py"
 
-
+    def make_env(self):
+        env = super().make_env()
+        env["XLA_PYTHON_CLIENT_PREALLOCATE"] = "False"
+        return env
+
 __pack__ = BraxBenchmark
diff --git a/benchmarks/brax/requirements.rocm.txt b/benchmarks/brax/requirements.rocm.txt