Clean up

Dockerfile

@@ -1,19 +1,12 @@
 # Stage 1: Base
 FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 as base
 
-ARG INDEX_URL
-ARG TORCH_VERSION
-ARG XFORMERS_VERSION
-
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 ENV DEBIAN_FRONTEND=noninteractive \
     TZ=Europe/London \
     PYTHONUNBUFFERED=1 \
     SHELL=/bin/bash
 
-# Create workspace working directory
-WORKDIR /
-
 # Install Ubuntu packages
 RUN apt update && \
     apt -y upgrade && \
@@ -69,6 +62,9 @@ RUN apt update && \
 RUN ln -s /usr/bin/python3.10 /usr/bin/python
 
 # Install Torch, xformers and tensorrt
+ARG INDEX_URL
+ARG TORCH_VERSION
+ARG XFORMERS_VERSION
 RUN pip3 install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url ${INDEX_URL} && \
     pip3 install --no-cache-dir xformers==${XFORMERS_VERSION} --index-url ${INDEX_URL} &&  \
     pip3 install --no-cache-dir tensorrt