diff --git a/docker/Dockerfile.nvidia b/docker/Dockerfile.nvidia
index d7e0a64d49..f303ad93c9 100644
--- a/docker/Dockerfile.nvidia
+++ b/docker/Dockerfile.nvidia
@@ -47,10 +47,6 @@ RUN apt-get update && apt-get install -y -q \
         nodejs ffmpeg gcc-offload-nvptx \
         texlive-latex-extra texlive-fonts-recommended dvipng cm-super
 
-# MPI_VER options 3,4,HPCX
-ARG MPI_VER=HPCX
-ENV MPIVER=${MPI_VER}
-
 # nvidia-container-runtime
 ENV NVIDIA_VISIBLE_DEVICES all
 ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
@@ -114,15 +110,8 @@ ENV LD_LIBRARY_PATH $HPCSDK_HOME/cuda/lib:$HPCSDK_HOME/cuda/lib64:$HPCSDK_HOME/c
 ENV CPATH $HPCSDK_HOME/comm_libs/mpi/include:$HPCSDK_HOME/comm_libs/nvshmem/include:$HPCSDK_HOME/comm_libs/nccl/include:${CPATH}
 
 # MPI
-RUN if [ "x$MPI_VER" = "x4" ]; then \
-        rm -f  $HPCSDK_HOME/comm_libs/mpi && \
-        ln -sf $HPCSDK_HOME/comm_libs/openmpi4/latest \
-               $HPCSDK_HOME/comm_libs/mpi ; \
-    elif [ "$MPI_VER" = "HPCX" ]; then \
-        rm -f  $HPCSDK_HOME/comm_libs/mpi && \
-        ln -sf $HPCSDK_HOME/comm_libs/hpcx/latest/ompi \
-               $HPCSDK_HOME/comm_libs/mpi ; \
-    fi;
+RUN rm -f  $HPCSDK_HOME/comm_libs/mpi && \
+    ln -sf $HPCSDK_HOME/comm_libs/hpcx/latest/ompi $HPCSDK_HOME/comm_libs/mpi  \
 
 # Install python nvidia dependencies
 RUN python3 -m venv /venv && \
@@ -148,7 +137,7 @@ FROM sdk-base as nvc
 ADD docker/nvdashboard.json /app/nvdashboard.json
 
 # mpi4py
-ENV MPI4PY_FLAGS='CC=nvc CFLAGS="-noswitcherror -tp=px"'
+ENV MPI4PY_FLAGS='$HPCSDK_HOME/comm_libs/hpcx/latest/hpcx-init.sh && hpcx_load && CC=nvc CFLAGS="-noswitcherror -tp=px"'
 
 ENV DEVITO_ARCH="nvc"
 ENV DEVITO_PLATFORM="nvidiaX"
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
index 077a84cd05..fe04ad5e70 100644
--- a/docker/entrypoint.sh
+++ b/docker/entrypoint.sh
@@ -2,7 +2,7 @@
 
 export PATH=/venv/bin:$PATH
 
-if [[ "$MPIVER" = "HPCX" ]]; then
+if [[ "$DEVITO_PLATFORM" = "nvidiaX" ]]; then
    echo "loading HPCX"
    source $HPCSDK_HOME/comm_libs/hpcx/latest/hpcx-init.sh
    hpcx_load