hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #3 from lucidrains/pytest-shard

Browse Source 
add pytest shard
This commit is contained in:
Phil Wang 2025-10-08 07:03:11 -07:00 committed by GitHub
commit 9c56ba0c9d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
.github/workflows/test.yml vendored
View File

@ -5,6 +5,11 @@ jobs:
build: build:
runs-on: ubuntu-latest runs-on: ubuntu-latest
timeout-minutes: 20
strategy:
fail-fast: false
matrix:
group: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
steps: steps:
- uses: actions/checkout@v4 - uses: actions/checkout@v4
@ -19,4 +24,4 @@ jobs:
python -m uv pip install -e .[test] python -m uv pip install -e .[test]
- name: Test with pytest - name: Test with pytest
run: | run: |
python -m pytest tests/ python -m pytest --num-shards 10 --shard-id ${{ matrix.group }} tests/

1
dreamer4/dreamer4.py
View File

@ -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 # 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_to_agent_token = add_reward_embed_to_agent_token
self.add_reward_embed_dropout = add_reward_embed_dropout
self.reward_encoder = SymExpTwoHot( self.reward_encoder = SymExpTwoHot(
**reward_encoder_kwargs, **reward_encoder_kwargs,

3
pyproject.toml
View File

@ -43,7 +43,8 @@ Repository = "https://github.com/lucidrains/dreamer4"
[project.optional-dependencies] [project.optional-dependencies]
examples = [] examples = []
test = [ test = [
"pytest" "pytest",
"pytest-shard"
] ]
[tool.pytest.ini_options] [tool.pytest.ini_options]