diff --git a/docker/Dockerfile b/docker/Dockerfile
index 01e51148..74c996ce 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -27,7 +27,8 @@ RUN apt remove cmake -y && \
 RUN pip install mpi4py
 
 # install onnx
-RUN pip install onnxruntime onnxruntime-gpu
+RUN pip install onnxruntime && \
+    if [ "$(uname -m)" = "x86_64" ]; then pip install onnxruntime-gpu; fi
 
 # hdf5 and h5py
 ENV HDF5_VERSION=2.0.0
@@ -74,8 +75,8 @@ RUN pip install git+https://github.com/NVIDIA/physicsnemo.git@v1.3.0
 
 # torch-harmonics
 RUN pip install setuptools_scm
-ENV FORCE_CUDA_EXTENSION 1
-ENV TORCH_CUDA_ARCH_LIST "8.0 8.6 8.7 9.0 10.0+PTX"
+ENV FORCE_CUDA_EXTENSION=1
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0a 10.0a+PTX"
 RUN cd /opt && git clone https://github.com/NVIDIA/torch-harmonics.git && \
     cd torch-harmonics && \
     pip install --no-build-isolation -e .
@@ -90,4 +91,4 @@ COPY makani /opt/makani/makani
 COPY tests /opt/makani/tests
 COPY pyproject.toml /opt/makani/pyproject.toml
 COPY README.md /opt/makani/README.md
-RUN cd /opt/makani && pip install --no-deps -e .
\ No newline at end of file
+RUN cd /opt/makani && pip install --no-deps -e .
diff --git a/docker/Dockerfile.aws b/docker/Dockerfile.aws
index c1549ffe..db099044 100644
--- a/docker/Dockerfile.aws
+++ b/docker/Dockerfile.aws
@@ -30,7 +30,8 @@ RUN apt remove cmake -y && \
 RUN pip install mpi4py
 
 # install onnx
-RUN pip install onnxruntime onnxruntime-gpu
+RUN pip install onnxruntime && \
+    if [ "$(uname -m)" = "x86_64" ]; then pip install onnxruntime-gpu; fi
 
 # install AWS prereqs for HDF5
 RUN cd /opt && git clone https://github.com/awslabs/aws-c-common.git && \
@@ -107,7 +108,7 @@ RUN cd /opt && git clone https://github.com/awslabs/aws-c-auth.git && \
     make install
 
 # 7. aws-c-s3 (final target for HDF5)
-RUn cd /opt && git clone https://github.com/awslabs/aws-c-s3.git && \
+RUN cd /opt && git clone https://github.com/awslabs/aws-c-s3.git && \
     cd aws-c-s3 && \
     mkdir build && cd build && \
     cmake .. -DCMAKE_BUILD_TYPE=Release && \
@@ -156,9 +157,8 @@ RUN pip install "git+https://github.com/google-research/weatherbench2.git"
 RUN pip install --ignore-installed "git+https://github.com/NVIDIA/mlperf-common.git"
 
 # torch-harmonics
-ENV FORCE_CUDA_EXTENSION 1
-ENV TORCH_CUDA_ARCH_LIST "8.0 8.6 8.7 9.0 10.0 12.0+PTX"
-ENV HARMONICS_VERSION 0.8.5
+ENV FORCE_CUDA_EXTENSION=1
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0a 10.0a 12.0+PTX"
 RUN cd /opt && git clone https://github.com/NVIDIA/torch-harmonics.git && \
     cd torch-harmonics && \
     pip install --no-build-isolation -e .
@@ -176,4 +176,4 @@ COPY makani /opt/makani/makani
 COPY tests /opt/makani/tests
 COPY pyproject.toml /opt/makani/pyproject.toml
 COPY README.md /opt/makani/README.md
-RUN cd /opt/makani && pip install -e .
\ No newline at end of file
+RUN cd /opt/makani && pip install -e .