From 10504cb00ed1356f50e2e162abe4eb8a4de0e999 Mon Sep 17 00:00:00 2001
From: Ramesh Sampath <1437573+sampathweb@users.noreply.github.com>
Date: Fri, 5 Jan 2024 06:52:04 +0530
Subject: [PATCH] Align keras requirments fo torch (#2272)

---
 requirements-tensorflow-cuda.txt | 1 -
 requirements-torch-cuda.txt      | 6 +++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/requirements-tensorflow-cuda.txt b/requirements-tensorflow-cuda.txt
index 4f474709fc..cacc9bdb7f 100644
--- a/requirements-tensorflow-cuda.txt
+++ b/requirements-tensorflow-cuda.txt
@@ -1,5 +1,4 @@
 # Tensorflow with cuda support.
---extra-index-url https://pypi.nvidia.com
 tf-nightly[and-cuda]==2.16.0.dev20231221  # Pin a working nightly until rc0.
 
 # Torch cpu-only version.
diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
index 126cbf0c4c..1ab7034525 100644
--- a/requirements-torch-cuda.txt
+++ b/requirements-torch-cuda.txt
@@ -2,9 +2,9 @@
 tf-nightly-cpu==2.16.0.dev20231221  # Pin a working nightly until rc0.
 
 # Torch with cuda support.
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.1.2+cu118
-torchvision==0.16.2+cu118
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.1.2
+torchvision==0.16.2
 
 # Jax cpu-only version.
 jax[cpu]