diff options
| author | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2022-03-22 20:51:32 +0100 |
|---|---|---|
| committer | Gustaf Rydholm <gustaf.rydholm@gmail.com> | 2022-03-22 20:51:32 +0100 |
| commit | dcdb2f4e11962b5f82c184e4deb1b9c0d51fdf95 (patch) | |
| tree | d23d0a3c54837f5bc1e6d739847898bc5847d2fc /pyproject.toml | |
| parent | e8663d381446c8c1024ac92af2760ccfe917d76b (diff) | |
build: update to jax with gpu support
Diffstat (limited to 'pyproject.toml')
| -rw-r--r-- | pyproject.toml | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/pyproject.toml b/pyproject.toml index d3ea39d..9a2eed3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,9 +6,9 @@ authors = ["Gustaf Rydholm <gustaf.rydholm@gmail.com>"] [tool.poetry.dependencies] python = "^3.9" -jax = "^0.3.4" jupyter = "^1.0.0" -jaxlib = "^0.3.2" +jaxlib = {url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.2+cuda11.cudnn82-cp39-none-manylinux2010_x86_64.whl"} +jax = "^0.3.4" [tool.poetry.dev-dependencies] pytest = "^5.2" |