fix case with no bias but also no causal_mask

openxla · Jan 12, 2024 · 3969975 · 3969975
1 parent 84e9ff3
commit 3969975
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xla/stream_executor/cuda/cuda_dnn.cc b/xla/stream_executor/cuda/cuda_dnn.cc
@@ -9305,7 +9305,7 @@ CudnnSupport::FusedMHARunnerFromDesc(
     scalar_input_values.push_back(dropout_scale);
     dropout_rng_offset = GetDropoutRngOffset(intermediate_shape);
 
-    if (bias_descriptor == std::nullopt) {
+    if (is_causal_mask) {
       // push negative infinity here
       scalar_input_uids.push_back(CudnnfMHAUid::NEG_INFINITY_ID);
       double negative_infinity_value = -std::numeric_limits<float>::infinity();