Skip to content

Commit 65dfd46

Browse files
Cristian GarciaFlax Authors
Cristian Garcia
authored and
Flax Authors
committed
MultiHeadAttention only keeps rngs if dropout_rate is positive
PiperOrigin-RevId: 759468680
1 parent 45ce6ee commit 65dfd46

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

flax/nnx/nn/attention.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -427,7 +427,7 @@ def __init__(
427427
dot_general_cls=self.out_dot_general_cls,
428428
rngs=rngs,
429429
)
430-
self.rngs = rngs if keep_rngs else None
430+
self.rngs = rngs if keep_rngs and dropout_rate > 0 else None
431431

432432
self.cached_key: nnx.Cache[Array] | None = None
433433
self.cached_value: nnx.Cache[Array] | None = None

0 commit comments

Comments
 (0)