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/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,
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]