add test to ensure correct torch version is installed

README.md

@@ -34,21 +34,21 @@ Tags follow these patterns:
 ##### _CUDA_
 - `:[pytorch-version]-py[python-version]-cuda-[x.x.x]-base-[ubuntu-version]`
 
-- `:latest-cuda` → `:2.1.1-py3.12-cuda-12.2.0-base-22.04`
+- `:latest-cuda` → `:2.1.1-py3.11-cuda-12.1.0-base-22.04`
 
 ##### _ROCm_
 - `:[pytorch-version]-py[python-version]-rocm-[x.x.x]-runtime-[ubuntu-version]`
 
-- `:latest-rocm` → `:2.1.1-py3.12-rocm-5.6-runtime-22.04`
+- `:latest-rocm` → `:2.1.1-py3.11-rocm-5.6-runtime-22.04`
 
 ##### _CPU_
 - `:[pytorch-version]-py[python-version]-ubuntu-[ubuntu-version]`
 
-- `:latest-cpu` → `:2.1.1-py3.12-cpu-22.04` 
+- `:latest-cpu` → `:2.1.1-py3.11-cpu-22.04` 
 
 Browse [here](https://github.com/ai-dock/jupyter-pytorch/pkgs/container/jupyter-pytorch) for an image suitable for your target environment.
 
-Supported Python versions: `3.12`, `3.11`, `3.10`
+Supported Python versions: `3.11`, `3.10`
 
 Supported Pytorch versions: `2.1.1`, `2.1.0` `2.0.1`
 

build/COPY_ROOT/opt/ai-dock/bin/build/layer0/init.sh

@@ -16,4 +16,6 @@ else
     exit 1
 fi
 
+$MAMBA_DEFAULT_RUN python /opt/ai-dock/tests/assert-torch-version.py
+
 /opt/ai-dock/bin/build/layer0/clean.sh