Torch 2.2 upgrade - Part 1 (#976)

mosaicml · Feb 15, 2024 · 1ef7409 · 1ef7409
1 parent 2e8982e
commit 1ef7409
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
@@ -23,6 +23,12 @@ jobs:
         - name: "2.1.2_cu121_flash2_aws"
           base_image: mosaicml/pytorch:2.1.2_cu121-python3.10-ubuntu20.04-aws
           dep_groups: "[gpu-flash2]"
+        - name: "2.2.0_cu121_flash2"
+          base_image: mosaicml/pytorch:2.2.0_cu121-python3.11-ubuntu20.04
+          dep_groups: "[gpu-flash2]"
+        - name: "2.2.0_cu121_flash2_aws"
+          base_image: mosaicml/pytorch:2.2.0_cu121-python3.11-ubuntu20.04-aws
+          dep_groups: "[gpu-flash2]"
     steps:
     - name: Maximize Build Space on Worker
       uses: easimon/maximize-build-space@v4

diff --git a/setup.py b/setup.py
@@ -55,7 +55,7 @@
     'accelerate>=0.25,<0.26',  # for HF inference `device_map`
     'transformers>=4.37,<4.38',
     'mosaicml-streaming>=0.7.4,<0.8',
-    'torch>=2.1,<2.2',
+    'torch>=2.1,<2.3',
     'datasets>=2.16,<2.17',
     'fsspec==2023.6.0',  # newer version results in a bug in datasets that duplicates data
     'sentencepiece==0.1.97',