tangger/openpi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

remove threefry setting in jax 0.5.0 (#279)

Browse Source 
continuation conversation from: https://app.graphite.dev/github/pr/Physical-Intelligence/monopi/6672/upgrade-jax-to-0-5-0?utm_source=gt-slack-notif&panel=timeline#comment-PRRC_kwDOLnRTkc50DYoL
This commit is contained in:
Haohuan Wang
2025-02-08 09:52:28 -08:00
committed by GitHub
parent 007e2b91ed ed05e55074
commit b7c8bf24d4
1 changed files with 0 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
scripts/train.py
View File

@@ -199,7 +199,6 @@ def main(config: _config.TrainConfig):
f"Batch size {config.batch_size} must be divisible by the number of devices {jax.device_count()}."
)
jax.config.update("jax_threefry_partitionable", True) # noqa: FBT003
jax.config.update("jax_compilation_cache_dir", str(epath.Path("~/.cache/jax").expanduser()))
rng = jax.random.key(config.seed)