From 63b63dfeddd6e21c0ed374294c1be5aa00989ed3 Mon Sep 17 00:00:00 2001 From: lucidrains Date: Wed, 8 Oct 2025 06:56:03 -0700 Subject: [PATCH 1/2] add shard --- .github/workflows/test.yml | 7 ++++++- pyproject.toml | 3 ++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 477de16..04038af 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -5,6 +5,11 @@ jobs: build: runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + group: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] steps: - uses: actions/checkout@v4 @@ -19,4 +24,4 @@ jobs: python -m uv pip install -e .[test] - name: Test with pytest run: | - python -m pytest tests/ + python -m pytest --num-shards 10 --shard-id ${{ matrix.group }} tests/ diff --git a/pyproject.toml b/pyproject.toml index 8cc7f77..65e3543 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,8 @@ Repository = "https://github.com/lucidrains/dreamer4" [project.optional-dependencies] examples = [] test = [ - "pytest" + "pytest", + "pytest-shard" ] [tool.pytest.ini_options] From b5744237bf144b062bb13774834a6a387d7a9ac0 Mon Sep 17 00:00:00 2001 From: lucidrains Date: Wed, 8 Oct 2025 06:58:46 -0700 Subject: [PATCH 2/2] fix --- dreamer4/dreamer4.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py index b5c8cd0..1b2a274 100644 --- a/dreamer4/dreamer4.py +++ b/dreamer4/dreamer4.py @@ -1175,6 +1175,7 @@ class DynamicsModel(Module): # each agent token will have the reward embedding of the previous time step - but could eventually just give reward its own token self.add_reward_embed_to_agent_token = add_reward_embed_to_agent_token + self.add_reward_embed_dropout = add_reward_embed_dropout self.reward_encoder = SymExpTwoHot( **reward_encoder_kwargs,