diff --git a/Dockerfile b/Dockerfile
index 3252107..25a002d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,6 +3,6 @@ FROM git.sandbox.iuk.hdm-stuttgart.de/grosse/jupyterlab-datascience:4.0.12d
 ARG NB_USER="jovyan"
 ARG NB_UID="1000"
 ARG NB_GID="100"
-ENV CONDA_CUDA_OVERRIDE="11.2"  
+ENV CONDA_CUDA_OVERRIDE="12.1"  
 ENV NUM_GPUS=1
-RUN CONDA_OVERRIDE_CUDA="11.2" conda install -c conda-forge pytorch-gpu
\ No newline at end of file
+RUN CONDA_OVERRIDE_CUDA="12.1" conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
\ No newline at end of file