diff --git a/container.def b/container.def
index 204172ca6dc6acd80aad8220f385d8fee9cfc703..2eaedb0befc229f98db885ef72dd0eaadf11ba86 100644
--- a/container.def
+++ b/container.def
@@ -1,23 +1,31 @@
 # Get the base image from nvidia/cuda Docker image.
-BootStrap: docker
-From: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 
+Bootstrap: docker
+From: docker.io/pytorch/pytorch:2.1.1-cuda12.1-cudnn8-runtime
 
 %post
     # Install special software or additions
     apt-get update -y # Good practice, update the package database.
-    apt-get install -y      libopenblas-dev # Install the BLAS.
+    apt-get install -y   git 
     apt-get autoremove -y && apt-get autoclean -y # Good practice, to keep the Docker image as small as possible.
 
+    conda update -y conda
+    conda install -y cupy compilers pkg-config "libjpeg-turbo>=3.0.0" opencv numba -c pytorch -c conda-forge
+    conda install -y pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
+    conda clean -a
+
+    pip install timm wandb imageio pandas gin-config tqdm scipy
+    pip cache purge
+
 %environment
     # Set some environment vars
-    export PATH=/opt/root/bin:$PATH
-    export LD_LIBRARY_PATH=/opt/root/lib:$LD_LIBRARY_PATH
-    export PYTHONPATH=/opt/root/lib
+    #export PYTHONPATH=/opt/root/lib
 
 
 %runscript
     # The command that will run when container is started via apptainer run command.
-    python /opt/root/tutorials/roofit/rf101_basics.py
+    echo "Container was created $NOW"
+    echo "Arguments received: $*"
+    $@
 
 %labels
     # Add custom metadata to the container.