Updated README.md and pinned packages

mdsunivie · Sep 21, 2023 · 11d3160 · 11d3160
1 parent 19c49e4
commit 11d3160
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -40,7 +40,7 @@ This will install the repository "in-place", so you can make changes to the sour
 If you need CUDA support to run the JAX code on GPUs (recommended), additionally install the prepackaged jax[cuda] wheel:
 
 ```bash
- pip install --upgrade jax==0.3.23 jaxlib==0.3.22+cuda11.cudnn82 dm-haiku==0.0.9 flax==0.6.4 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+ pip install --upgrade jax==0.3.23 jaxlib==0.3.22+cuda11.cudnn82 dm-haiku==0.0.9 flax==0.6.4 orbax==0.1.2 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 

diff --git a/setup.cfg b/setup.cfg
@@ -24,14 +24,15 @@ setup_requires =
 install_requires =
  jax==0.3.23
  jaxlib==0.3.22
+ orbax-checkpoint==0.1.1
  dm-haiku
  optax
  chex==0.1.5
  e3nn-jax==0.17.1
  pyscf
  h5py==3.2.1
  ruamel.yaml
- pydantic
+ pydantic==1.10.11
  wandb
  pandas
  ase