diff --git a/container.def b/container.def
index 204172ca6dc6acd80aad8220f385d8fee9cfc703..65c850779288aa700d913c3a978f55f26893df57 100644
--- a/container.def
+++ b/container.def
@@ -5,19 +5,15 @@ From: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
 %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 autoremove -y && apt-get autoclean -y # Good practice, to keep the Docker image as small as possible.
-
+    apt-get install -y \
+        wget curl vim bzip2 git build-essential \
+        libopenblas-dev 
+    apt-get autoremove -y && apt-get clean && rm -rf /var/lib/apt/lists/* # keep the Docker image as small as possible.
+    
 %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
-
-
-%runscript
-    # The command that will run when container is started via apptainer run command.
-    python /opt/root/tutorials/roofit/rf101_basics.py
+    export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6+PTX"
+    export TCNN_CUDA_ARCHITECTURES="75;80;86"
 
 %labels
     # Add custom metadata to the container.