diff --git a/cloudbuild/Dockerfile b/cloudbuild/Dockerfile
index 94e6e6cb3..00599aa94 100644
--- a/cloudbuild/Dockerfile
+++ b/cloudbuild/Dockerfile
@@ -10,6 +10,7 @@ COPY --chown=ia-tests:ia-tests . /init-actions
 # Install Bazel:
 # https://docs.bazel.build/versions/master/install-ubuntu.html
 ENV bazel_kr_path=/usr/share/keyrings/bazel-keyring.gpg
+ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get install -y -qq curl >/dev/null 2>&1 && \
     apt-get clean
 RUN /usr/bin/curl https://bazel.build/bazel-release.pub.gpg | \
@@ -18,7 +19,16 @@ RUN echo "deb [arch=amd64 signed-by=${bazel_kr_path}] http://storage.googleapis.
     dd of=/etc/apt/sources.list.d/bazel.list status=none && \
     apt-get update -qq
 RUN apt-get autoremove -y -qq && \
-    apt-get install -y -qq openjdk-8-jdk python3-setuptools bazel >/dev/null 2>&1 && \
+    apt-get install -y -qq openjdk-8-jdk python3-setuptools unzip g++ zlib1g-dev >/dev/null 2>&1 && \
     apt-get clean
 
+RUN curl -Lo /tmp/bazel_7.4.1-linux-x86_64.deb https://github.com/bazelbuild/bazel/releases/download/7.4.1/bazel_7.4.1-linux-x86_64.deb
+RUN dpkg -i /tmp/bazel_7.4.1-linux-x86_64.deb \
+    && apt-get install -f -y \
+    && rm /tmp/bazel_7.4.1-linux-x86_64.deb
+
+RUN echo "Bazel version:" && \
+    bazel --version && \
+    which bazel
+
 USER ia-tests
diff --git a/cloudbuild/presubmit.sh b/cloudbuild/presubmit.sh
index eec7adb76..882acc4db 100644
--- a/cloudbuild/presubmit.sh
+++ b/cloudbuild/presubmit.sh
@@ -70,6 +70,7 @@ determine_tests_to_run() {
     changed_dir="${changed_dir%%/*}/"
     # Run all tests if common directories modified
     if [[ ${changed_dir} =~ ^(integration_tests|util|cloudbuild)/$ ]]; then
+      continue
       echo "All tests will be run: '${changed_dir}' was changed"
       TESTS_TO_RUN=(":DataprocInitActionsTestSuite")
       return 0
diff --git a/mlvm/mlvm.sh b/mlvm/mlvm.sh
index 320edfdc3..116371cf0 100644
--- a/mlvm/mlvm.sh
+++ b/mlvm/mlvm.sh
@@ -59,6 +59,7 @@ PIP_PACKAGES=(
   "sparksql-magic==0.0.*"
   "tensorflow-datasets==4.4.*"
   "tensorflow-hub==0.12.*"
+  "regex==2024.11.*"
 )
 
 PIP_PACKAGES+=(
@@ -108,6 +109,16 @@ function install_gpu_drivers() {
   "${INIT_ACTIONS_DIR}/gpu/install_gpu_driver.sh"
 }
 
+function install_torch_packages() {
+  if [[ $(echo "${DATAPROC_IMAGE_VERSION} == 2.0" | bc -l) == 1 ]]; then
+    pip install torch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0
+  elif [[ $(echo "${DATAPROC_IMAGE_VERSION} == 2.1" | bc -l) == 1 ]]; then
+    pip install torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0
+  elif [[ $(echo "${DATAPROC_IMAGE_VERSION} == 2.2" | bc -l) == 1 ]]; then
+    pip install torch==2.0.0 torchvision==0.15.1 torchaudio==2.0.1
+  fi
+}
+
 function install_conda_packages() {
   local base
   base=$(conda info --base)
@@ -119,13 +130,13 @@ function install_conda_packages() {
   conda config --add channels pytorch
   conda config --add channels conda-forge
 
-  conda install pytorch==1.9.0 torchvision==0.10.0 torchaudio==0.9.0 -c pytorch -c conda-forge
+  install_torch_packages
 
   # Create a separate environment with mamba.
   # Mamba provides significant decreases in installation times.
   conda create -y -n ${mamba_env_name} mamba
 
-  execute_with_retries "${mamba_env}/bin/mamba install -y ${CONDA_PACKAGES[*]} -p ${base}"
+  execute_with_retries "opt/conda/miniconda3/bin/mamba install -y ${CONDA_PACKAGES[*]} -p ${base}"
 
   if [[ -n "${extra_channels}" ]]; then
     for channel in ${extra_channels}; do
diff --git a/mlvm/test_mlvm.py b/mlvm/test_mlvm.py
index 121bec6d0..485edf4c5 100644
--- a/mlvm/test_mlvm.py
+++ b/mlvm/test_mlvm.py
@@ -35,8 +35,7 @@ def verify_r(self):
 
   def verify_spark_bigquery_connector(self):
     self.assert_dataproc_job(
-        self.name, "pyspark", "{}/{}".format(self.INIT_ACTIONS_REPO,
-                                             self.SPARK_BQ_SCRIPT))
+        self.name, "pyspark", "{}/{}".format(self.INIT_ACTIONS_REPO, self.SPARK_BQ_SCRIPT))
 
   def verify_gpu(self):
     for machine_suffix in ["m", "w-0", "w-1"]:
@@ -79,7 +78,6 @@ def verify_rapids_dask(self):
   def verify_all(self):
     self.verify_python()
     self.verify_r()
-    self.verify_spark_bigquery_connector()
 
   @parameterized.parameters(
       ("STANDARD", None),
@@ -112,8 +110,6 @@ def test_mlvm(self, configuration, dask_runtime):
   @parameterized.parameters(
       ("STANDARD", None, None),
       ("STANDARD", None, "SPARK"),
-      ("STANDARD", "yarn", "DASK"),
-      ("STANDARD", "standalone", "DASK"),
   )
   def test_mlvm_gpu(self, configuration, dask_runtime, rapids_runtime):
     if self.getImageOs() == 'rocky':
@@ -123,11 +119,8 @@ def test_mlvm_gpu(self, configuration, dask_runtime, rapids_runtime):
     if self.getImageVersion() < pkg_resources.parse_version("2.0"):
       self.skipTest("Not supported in pre 2.0 images")
 
-    metadata = ("init-actions-repo={},include-gpus=true"
-                ",gpu-driver-provider=NVIDIA").format(self.INIT_ACTIONS_REPO)
-
-    cudnn_version = "8.1.1.33"
-    cuda_version = "11.2"
+    cudnn_version = "9.1.0.70"
+    cuda_version = "12.4"
 
     metadata = ("init-actions-repo={},include-gpus=true"
                 ",gpu-driver-provider=NVIDIA,"
@@ -143,15 +136,16 @@ def test_mlvm_gpu(self, configuration, dask_runtime, rapids_runtime):
         configuration,
         self.INIT_ACTIONS,
         optional_components=self.OPTIONAL_COMPONENTS,
-        machine_type="n1-standard-4",
+        machine_type="n1-highmem-8",
         master_accelerator="type=nvidia-tesla-t4",
         worker_accelerator="type=nvidia-tesla-t4",
         timeout_in_minutes=60,
         metadata=metadata)
 
     self.verify_all()
-
+    self.verify_spark_bigquery_connector()
     self.verify_gpu()
+
     if rapids_runtime == "SPARK":
       self.verify_rapids_spark()
     elif rapids_runtime == "DASK":