diff --git a/setup_flax.sh b/setup_flax.sh
index ee465bd..0376eed 100644
--- a/setup_flax.sh
+++ b/setup_flax.sh
@@ -2,7 +2,7 @@
 
 set -e
 
-pip install -r requirements.txt
+pip install torch flax wandb
 
 # download vqgan
 mkdir -p pretrained/vqgan