From 5c1887555a05bbd3489960eda0ca9d92ab4fdbdb Mon Sep 17 00:00:00 2001
From: magpowell <mp4257@columbia.edu>
Date: Tue, 7 May 2024 13:52:06 -0400
Subject: [PATCH] expand requirements

---
 requirements.txt | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 5ad8cf8..449726d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,10 +2,11 @@ dinosaur@git+https://github.com/google-research/dinosaur
 dm-haiku
 gin-config
 jax
-jaxlib
+jaxlib==0.1.71+cuda111
 optax
 tensorflow_probability
 jupyterlab==4.1.6 
 gcsfs==2024.3.1 
 matplotlib==3.8.4 
 julia==0.6.2
+nvidia-cudnn-cu11