diff --git a/people-and-planet-ai/weather-forecasting/notebooks/3-training.ipynb b/people-and-planet-ai/weather-forecasting/notebooks/3-training.ipynb
index 56be23f2fd3..f0656c1208c 100644
--- a/people-and-planet-ai/weather-forecasting/notebooks/3-training.ipynb
+++ b/people-and-planet-ai/weather-forecasting/notebooks/3-training.ipynb
@@ -1381,7 +1381,7 @@
     "    display_name=\"weather-forecasting\",\n",
     "    python_package_gcs_uri=f\"gs://{bucket}/weather/weather-model-1.0.0.tar.gz\",\n",
     "    python_module_name=\"weather.trainer\",\n",
-    "    container_uri=\"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.2-2.py310:latest\",\n",
+    "    container_uri=\"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.2-4.py310:latest\",\n",
     ")\n",
     "job.run(\n",
     "    machine_type=\"n1-highmem-8\",\n",
diff --git a/people-and-planet-ai/weather-forecasting/serving/weather-model/pyproject.toml b/people-and-planet-ai/weather-forecasting/serving/weather-model/pyproject.toml
index e016d2061c9..e5b3a98ffbe 100644
--- a/people-and-planet-ai/weather-forecasting/serving/weather-model/pyproject.toml
+++ b/people-and-planet-ai/weather-forecasting/serving/weather-model/pyproject.toml
@@ -18,7 +18,7 @@ name = "weather-model"
 version = "1.0.0"
 dependencies = [
     "datasets==3.0.1",
-    "torch==2.2.0",  # make sure this matches the `container_uri` in `notebooks/3-training.ipynb`
+    "torch==2.4.0",  # make sure this matches the `container_uri` in `notebooks/3-training.ipynb`
     "transformers==4.48.0",
 ]