diff --git a/ai/containers/Dockerfile b/ai/containers/Dockerfile
deleted file mode 100644
index dd2364f6..00000000
--- a/ai/containers/Dockerfile
+++ /dev/null
@@ -1,33 +0,0 @@
-# Set the base image to Python 3.9.17-slim
-FROM python:3.9.17-slim AS build
-
-
-WORKDIR ai
-
-# Copy the required files to the container
-COPY env-files  env-files
-COPY src        src
-COPY tests      tests
-COPY setup.py   setup.py
-COPY README.md  README.md
-
-# Set environment variable
-ENV LC_ALL=C
-
-# Install packages
-RUN apt update && apt install -y iputils-ping && apt install -y telnet && apt install -y krb5-user && apt install -y nano
-RUN apt-get clean && rm -rf /var/lib/apt/lists/*
-
-# Update pip and install python requirements
-RUN pip install --upgrade pip
-RUN pip install -r /ai/env-files/pytorch-env-container.txt
-RUN pip install -e /ai
-
-# Define the run command when the container starts
-CMD ["sleep", "300"]
-CMD ["python", "-V"]
-
-# Additional information labels
-LABEL Author="azoechba"
-LABEL Version="v0.0.1"
-LABEL Description="This is the base image of the itwinai library in Docker."
diff --git a/ai/env-files/pytorch-env-container.txt b/ai/env-files/pytorch-env-container.txt
deleted file mode 100644
index ee5eb5e4..00000000
--- a/ai/env-files/pytorch-env-container.txt
+++ /dev/null
@@ -1,11 +0,0 @@
-torch==1.13.1
-torchvision==0.14.1
-lightning==2.0.0
-torchmetrics
-mlflow>=2
-typer
-rich
-pyyaml
-omegaconf
-typing-extensions==4.5.0
-typing_extensions==4.5.0
\ No newline at end of file
diff --git a/containers/Dockerfile b/containers/Dockerfile
new file mode 100644
index 00000000..8e8b4b60
--- /dev/null
+++ b/containers/Dockerfile
@@ -0,0 +1,42 @@
+# Set the base image to Python 3.9.17-slim
+#FROM python:3.9.17-slim AS build
+
+
+#Create from the following container
+FROM nvcr.io/nvidia/pytorch:21.12-py3 AS build
+#docker pull nvcr.io/nvidia/pytorch:21.12-py3
+
+
+
+WORKDIR ai
+
+# Copy the required files to the container
+COPY env-files      env-files
+COPY src            src
+COPY tests          tests
+COPY pyproject.toml pyproject.toml 
+#COPY setup.py   setup.py
+#COPY README.md  README.md
+
+# Set environment variable
+ENV LC_ALL=C
+
+# Install packages
+#RUN apt update && apt install -y iputils-ping && apt install -y telnet && apt install -y krb5-user && apt install -y nano
+#RUN apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Update pip and install python requirements
+RUN pip install --upgrade pip
+#RUN pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
+RUN ls 
+RUN pip install -r env-files/torch/pytorch-env-gpu-container.txt
+RUN pip install -e .
+
+# Define the run command when the container starts
+CMD ["sleep", "300"]
+CMD ["python", "-V"]
+
+# Additional information labels
+LABEL Author="azoechba"
+LABEL Version="v0.0.1"
+LABEL Description="This is the base image of the itwinai library in Docker."
diff --git a/ai/containers/base-container.def b/containers/base-container.def
similarity index 100%
rename from ai/containers/base-container.def
rename to containers/base-container.def
diff --git a/ai/containers/build-base-container.sh b/containers/build-base-container.sh
similarity index 100%
rename from ai/containers/build-base-container.sh
rename to containers/build-base-container.sh
diff --git a/ai/containers/build-train-container.sh b/containers/build-train-container.sh
similarity index 100%
rename from ai/containers/build-train-container.sh
rename to containers/build-train-container.sh
diff --git a/ai/containers/dummy-deployment.yaml b/containers/dummy-deployment.yaml
similarity index 100%
rename from ai/containers/dummy-deployment.yaml
rename to containers/dummy-deployment.yaml
diff --git a/ai/containers/dummy-pod-jsc.yaml b/containers/dummy-pod-jsc.yaml
similarity index 100%
rename from ai/containers/dummy-pod-jsc.yaml
rename to containers/dummy-pod-jsc.yaml
diff --git a/ai/containers/eosclient-deployment.yaml b/containers/eosclient-deployment.yaml
similarity index 100%
rename from ai/containers/eosclient-deployment.yaml
rename to containers/eosclient-deployment.yaml
diff --git a/ai/containers/intertwin-pod.yaml b/containers/intertwin-pod.yaml
similarity index 100%
rename from ai/containers/intertwin-pod.yaml
rename to containers/intertwin-pod.yaml
diff --git a/ai/containers/train-container.def b/containers/train-container.def
similarity index 100%
rename from ai/containers/train-container.def
rename to containers/train-container.def
diff --git a/ai/containers/training-deployment-cern.yaml b/containers/training-deployment-cern.yaml
similarity index 100%
rename from ai/containers/training-deployment-cern.yaml
rename to containers/training-deployment-cern.yaml
diff --git a/ai/containers/training-pod-cern.yaml b/containers/training-pod-cern.yaml
similarity index 100%
rename from ai/containers/training-pod-cern.yaml
rename to containers/training-pod-cern.yaml
diff --git a/ai/containers/training-pod.yaml b/containers/training-pod.yaml
similarity index 100%
rename from ai/containers/training-pod.yaml
rename to containers/training-pod.yaml
diff --git a/env-files/torch/pytorch-env-gpu-container.txt b/env-files/torch/pytorch-env-gpu-container.txt
new file mode 100644
index 00000000..32aedc18
--- /dev/null
+++ b/env-files/torch/pytorch-env-gpu-container.txt
@@ -0,0 +1,12 @@
+mlflow>=2
+typer
+rich
+pyyaml
+omegaconf
+typing-extensions==4.5.0
+typing_extensions==4.5.0
+
+
+
+
+
diff --git a/pyproject.toml b/pyproject.toml
index fd5c42f0..b388c182 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ name = "itwinai"
 version = "0.1"
 description = "AI and ML workflows module for interTwin"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.8"
 license = { file = "LICENSE" }
 keywords = ["ml", "ai", "hpc"]
 authors = [
diff --git a/use-cases/mnist/env-files/preproc-env-container.txt b/use-cases/mnist/env-files/preproc-env-container.txt
new file mode 100644
index 00000000..feccf8ae
--- /dev/null
+++ b/use-cases/mnist/env-files/preproc-env-container.txt
@@ -0,0 +1,3 @@
+torch
+torchvision
+