diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index c563fa27..3e94f730 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -202,4 +202,5 @@ jobs: export EXTRA_PYTEST="${{ needs.build-and-push.outputs.extra_pytest }}" export HF_TOKEN=${{ secrets.HF_TOKEN }} echo $DOCKER_IMAGE + docker pull $DOCKER_IMAGE pytest -s -vv integration-tests ${PYTEST_FLAGS} ${EXTRA_PYTEST} diff --git a/integration-tests/models/test_bloom_560m.py b/integration-tests/models/test_bloom_560m.py index d413519e..c7859181 100644 --- a/integration-tests/models/test_bloom_560m.py +++ b/integration-tests/models/test_bloom_560m.py @@ -3,7 +3,7 @@ import pytest @pytest.fixture(scope="module") def bloom_560_handle(launcher): - with launcher("bigscience/bloom-560m") as handle: + with launcher("bigscience/bloom-560m", num_shard=1) as handle: yield handle