Update Dockerfile

Dockerfile

@@ -43,11 +43,8 @@ RUN apt-get update && apt-get install -y \
 
 ENV PATH =/opt/aws/neuron/bin:$PATH
 
-ARG OPTIMUM_NEURON_BRANCH=main
-RUN git clone --branch $OPTIMUM_NEURON_BRANCH https://github.com/huggingface/optimum-neuron.git /tmp/optimum-neuron && \
-    pip install --no-cache-dir "/tmp/optimum-neuron[neuronx]" && \
-    rm -rf /tmp/optimum-neuron
-
+ARG OPTIMUM_NEURON_VERSION=0.4.0
+RUN pip install --extra-index-url https://pip.repos.neuron.amazonaws.com --no-cache-dir optimum-neuron[neuronx]==$OPTIMUM_NEURON_VERSION
 # End Neuron setup
 RUN pip install --no-cache-dir pip -U && \
 	pip install --no-cache-dir \
@@ -89,4 +86,4 @@ ENV PYTHONPATH=$HOME/app \
 	GRADIO_THEME=huggingface \
 	TQDM_POSITION=-1 \
 	TQDM_MININTERVAL=1 \
-	SYSTEM=spaces 
+	SYSTEM=spaces