add contiguous to scale too

keras-team · Nov 7, 2024 · 3e34498 · 3e34498
1 parent a36d119
commit 3e34498
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/keras/src/backend/torch/nn.py b/keras/src/backend/torch/nn.py
@@ -956,6 +956,8 @@ def dot_product_attention(
     else:
         if mask is not None:
             mask = mask.contiguous()
+        if scale is not None:
+            scale = scale.contiguous()
         attention_output = torch.nn.functional.scaled_dot_product_attention(
             query.contiguous(),
             key.contiguous(),