From cd8372101f9244b0f854cff1c328e6a5581e39be Mon Sep 17 00:00:00 2001
From: Ashley Kleynhans <ashley.kleynhans@gmail.com>
Date: Tue, 26 Mar 2024 17:18:40 +0200
Subject: [PATCH] Clean up

---
 Dockerfile | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 048add7..7eaffe3 100644
--- a/Dockerfile
+++ b/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