diff --git a/keras_hub/src/models/pali_gemma/pali_gemma_vit.py b/keras_hub/src/models/pali_gemma/pali_gemma_vit.py
index b217509541..621b4562de 100644
--- a/keras_hub/src/models/pali_gemma/pali_gemma_vit.py
+++ b/keras_hub/src/models/pali_gemma/pali_gemma_vit.py
@@ -204,9 +204,8 @@ def __init__(
         self.intermediate_dim = intermediate_dim
 
     def compute_attention(self, x, mask=None):
-        mask = None
         if mask is not None:
-            mask = ops.cast(mask, dtype=x.dtype) if mask is not None else None
+            mask = ops.cast(mask, dtype=x.dtype)
         return self.attn(x, attention_mask=mask)[0]
 
     def build(self, input_shape):
diff --git a/requirements-torch-cuda.txt b/requirements-torch-cuda.txt
index 2a601f0f20..a696dc2a85 100644
--- a/requirements-torch-cuda.txt
+++ b/requirements-torch-cuda.txt
@@ -4,8 +4,8 @@ tensorflow-text~=2.18
 
 # Torch with cuda support.
 --extra-index-url https://download.pytorch.org/whl/cu121
-torch==2.5.1+cu121
-torchvision==0.20.1+cu121
+torch==2.6.0+cpu
+torchvision==0.21.0+cpu
 
 # Jax cpu-only version.
 jax[cpu]