Update container to newer version rather than updating specific packages

microsoft · Apr 5, 2024 · 203f5b7 · 203f5b7
1 parent 93d3d6a
commit 203f5b7
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/.github/workflows/nv-a6000.yml b/.github/workflows/nv-a6000.yml
@@ -23,7 +23,7 @@ jobs:
  unit-tests:
  runs-on: [self-hosted, nvidia, a6000]
  container:
- image: nvcr.io/nvidia/pytorch:23.03-py3
+ image: nvcr.io/nvidia/pytorch:23.09-py3
  ports:
  - 80
  options: --gpus all --shm-size "8G"
@@ -47,7 +47,6 @@ jobs:
  - name: Install deepspeed
  run: |
  python -m pip install docutils==0.18.1 jinja2==3.0 urllib3==1.26.11 ninja
- python -m pip install thinc spacy confection --upgrade
  python -m pip install .[dev,1bit,autotuning,inf]
  ds_report
  - name: Python environment
@@ -57,8 +56,8 @@ jobs:
  run: |
  unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
  cd tests
- python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2' unit/ --torch_ver="2.0" --cuda_ver="12"
- python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2_ops' unit/ --torch_ver="2.0" --cuda_ver="12"
+ python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2' unit/ --torch_ver="2.1" --cuda_ver="12"
+ python -m pytest --color=yes --durations=0 --verbose -rF -m 'inference_v2_ops' unit/ --torch_ver="2.1" --cuda_ver="12"
  - name: MII unit tests
  run: |
  BRANCH="main"