diff --git a/Dockerfile b/Dockerfile
index 11b884c..57ca661 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -29,8 +29,6 @@ RUN     apt-get -q update && \
 
 
 
-
-
 COPY env.yml /tmp/env.yml
 
 RUN conda env create --file /tmp/env.yml && \
@@ -40,8 +38,8 @@ RUN conda env create --file /tmp/env.yml && \
    # CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)")) && \
    # echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib' > $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh && \
     python -m ipykernel install --name=${KERNEL} && \
-    # Maybe just with the Runtimes? 
-    conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge && \
+
+    # conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge && \
     fix-permissions $CONDA_DIR && \
     fix-permissions /home/$NB_USER
 
@@ -49,6 +47,7 @@ RUN conda env create --file /tmp/env.yml && \
 USER jovyan
 # other packages here... 
 
+
 # From the "original" 2022 version: 
 # USER    $NB_UID:$NB_GID
 ENV     PATH=${PATH}:/usr/local/cuda/bin
diff --git a/env.yml b/env.yml
index 22570a5..f9e9747 100644
--- a/env.yml
+++ b/env.yml
@@ -8,14 +8,17 @@ dependencies:
 - python=3.7
 #- python=3.9
 #- tensorflow=2.9.1=cuda112*
-#- cudatoolkit=11.8.0
-#- cudnn=8.1.0
+- cudatoolkit=11.1
+- cudnn=8.0.5
 - jaxlib=0.1.55
 - tensorboard
 - ipykernel
 - matplotlib
 - pip
 - pip:
+  - click
+  - requests
+  - tqdm
   - gdown
   - imageio-ffmpeg==0.4.3
   - jax==0.1.73