From 6fb618f8018761617475c803d62b50f851ba57d1 Mon Sep 17 00:00:00 2001
From: Ved Patwardhan <54766411+vedpatwardhan@users.noreply.github.com>
Date: Mon, 5 Feb 2024 14:40:05 +0530
Subject: [PATCH] fix: hardcoded the index url for installing torchvision on
 the gpu dockerfile

---
 docker/gpu_framework_directory.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/docker/gpu_framework_directory.py b/docker/gpu_framework_directory.py
index ba365c1efb467..3dee4b7bf7750 100644
--- a/docker/gpu_framework_directory.py
+++ b/docker/gpu_framework_directory.py
@@ -43,6 +43,13 @@ def install_pkg(path, pkg, base="fw/"):
             "  --no-cache-dir",
             shell=True,
         )
+        subprocess.run(
+            f"yes |pip3 install --upgrade torchvision --index-url"
+            f" https://download.pytorch.org/whl/cu121 --target"
+            f" {path} --default-timeout=100 --extra-index-url"
+            " --no-cache-dir",
+            shell=True,
+        )
     elif pkg.split("==")[0] if "==" in pkg else pkg == "jax":
         subprocess.run(
             f"yes |pip install --upgrade --target {path} 'jax[cuda12_pip]' -f"