chore: Add PyTorch 2.0 images to prediction first-party uris.

PiperOrigin-RevId: 551059631
googleapis · Jul 26, 2023 · 3bfd996 · 3bfd996
1 parent 2788ef3
commit 3bfd996
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/constants/prediction.py b/google/cloud/aiplatform/constants/prediction.py
@@ -166,6 +166,12 @@
 ]
 
 PYTORCH_CONTAINER_URIS = [
+    "us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.2-0:latest",
+    "europe-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.2-0:latest",
+    "asia-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.2-0:latest",
+    "us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.2-0:latest",
+    "europe-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.2-0:latest",
+    "asia-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.2-0:latest",
     "us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-13:latest",
     "europe-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-13:latest",
     "asia-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-13:latest",