revert scripts

mosaicml · Nov 20, 2024 · d39181c · d39181c
1 parent 8399498
commit d39181c
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/scripts/train/benchmarking/submit_benchmarks.py b/scripts/train/benchmarking/submit_benchmarks.py
@@ -47,7 +47,7 @@ def parse_args():
     parser.add_argument(
         '--image',
         type=str,
-        default='mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu22.04',
+        default='mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04',
     )
     parser.add_argument(
         '--git_branch',

diff --git a/scripts/train/benchmarking/sweep.sh b/scripts/train/benchmarking/sweep.sh
@@ -2,7 +2,7 @@
 
 PROJECT="tput"
 GIT_COMMIT="v0.0.4"
-IMAGE="mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu22.04"
+IMAGE="mosaicml/pytorch:2.1.0_cu121-python3.10-ubuntu20.04"
 CLUSTER_40GB= # TODO
 
 for PRECISION in fp8 bf16

diff --git a/setup.py b/setup.py
@@ -57,7 +57,7 @@
     'accelerate>=0.25,<1.2',  # for HF inference `device_map`
     'transformers>=4.43.2,<4.47',
     'mosaicml-streaming>=0.9.0,<0.10',
-    'torch>=2.5.1,<2.4.1',
+    'torch>=2.5.1,<2.5.2',
     'datasets>=2.20.0,<2.21',
     'fsspec==2023.6.0',  # newer version results in a bug in datasets that duplicates data
     'sentencepiece==0.2.0',