diff --git a/Dockerfile.tmpl b/Dockerfile.tmpl
index 6eab7ca0..b9bcf2c8 100644
--- a/Dockerfile.tmpl
+++ b/Dockerfile.tmpl
@@ -81,7 +81,7 @@ RUN conda config --add channels nvidia && \
 
 # b/232247930: uninstall pyarrow to avoid double installation with the GPU specific version.
 RUN pip uninstall -y pyarrow && \
-    conda install cudf=21.10 cuml=21.10 cudatoolkit=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION && \
+    conda install cudf=21.12.02 cuml=21.12 cudatoolkit=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION && \
     /tmp/clean-layer.sh
 {{ end }}
 
@@ -97,7 +97,7 @@ RUN conda install implicit && \
 # Install PyTorch
 {{ if eq .Accelerator "gpu" }}
 COPY --from=torch_whl /tmp/whl/*.whl /tmp/torch/
-RUN conda install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION} && \
+RUN conda install -c pytorch magma-cuda${MAGMA_CUDA_VERSION} && \
     pip install /tmp/torch/*.whl && \
     rm -rf /tmp/torch && \
     /tmp/clean-layer.sh
diff --git a/config.txt b/config.txt
index 226e6f9e..ee1de9c9 100644
--- a/config.txt
+++ b/config.txt
@@ -9,4 +9,5 @@ TORCHAUDIO_VERSION=0.11.0
 TORCHTEXT_VERSION=0.12.0
 TORCHVISION_VERSION=0.12.0
 CUDA_MAJOR_VERSION=11
-CUDA_MINOR_VERSION=0
+CUDA_MINOR_VERSION=4
+MAGMA_CUDA_VERSION=113