From dcdb2f4e11962b5f82c184e4deb1b9c0d51fdf95 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Tue, 22 Mar 2022 20:51:32 +0100
Subject: build: update to jax with gpu support

---
 pyproject.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'pyproject.toml')

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"
-- 
cgit v1.2.3-70-g09d2