Update Dockerfile

Dockerfile

@@ -43,10 +43,9 @@ RUN apt-get update && apt-get install -y \
 
 ENV PATH =/opt/aws/neuron/bin:$PATH
 
-RUN pip install --no-cache-dir --extra-index-url https://pip.repos.neuron.amazonaws.com \
-    git+https://github.com/huggingface/optimum-neuron.git@cpu-compilation
-RUN pip install --no-cache-dir --extra-index-url https://pip.repos.neuron.amazonaws.com \
-    optimum-neuron[neuronx]
+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 ".[neuronx]" /tmp/optimum-neuron
 
 # End Neuron setup
 RUN pip install --no-cache-dir pip -U && \