Fix Pytorch GPU test (#2087)

* Update requirements-torch-cuda.txt

* Update requirements-torch-cuda.txt

* Update requirements.txt

* Update requirements-torch-cuda.txt

Author: divyashreepathihalli

requirements-torch-cuda.txt

@@ -3,9 +3,9 @@ tensorflow-cpu~=2.18
 tensorflow-text~=2.18
 
 # Torch with cuda support.
---extra-index-url https://download.pytorch.org/whl/cu121
-torch==2.6.0+cpu
-torchvision==0.21.0+cpu
+--extra-index-url https://download.pytorch.org/whl/cu126
+torch==2.6.0+cu126
+torchvision==0.21.0+cu126
 
 # Jax cpu-only version.
 jax[cpu]

requirements.txt

@@ -4,7 +4,7 @@ tensorflow-text~=2.18
 
 # Torch.
 --extra-index-url https://download.pytorch.org/whl/cpu
-torch>=2.1.0
+torch>=2.6.0
 torchvision>=0.16.0
 
 # Jax.