Update Dockerfile

Dockerfile

@@ -43,16 +43,9 @@ RUN apt-get update && apt-get install -y \
 
 ENV PATH =/opt/aws/neuron/bin:$PATH
 
-# Define branch and repo (optional args, easy to override at build time)
-ARG OPTIMUM_NEURON_BRANCH=cpu-compilation
-ARG OPTIMUM_NEURON_REPO=https://github.com/huggingface/optimum-neuron.git
-
-# Install from your branch
-# install the branch
-RUN pip install --no-cache-dir "git+https://github.com/huggingface/optimum-neuron.git@cpu-compilation"
-# add the extras
-RUN pip install --no-cache-dir "optimum-neuron[neuronx]"
-
+ARG OPTIMUM_NEURON_VERSION=0.3.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 \
 	datasets \
@@ -93,4 +86,4 @@ ENV PYTHONPATH=$HOME/app \
 	GRADIO_THEME=huggingface \
 	TQDM_POSITION=-1 \
 	TQDM_MININTERVAL=1 \
-	SYSTEM=spaces 
+	SYSTEM=spaces