Update Dockerfile

Dockerfile

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