Unverified Commit 7af7d7ce authored by Siarhei Melnik's avatar Siarhei Melnik Committed by GitHub
Browse files

fix: replace asserts by error (#13894)

parent f099249c
......@@ -172,7 +172,8 @@ class FlaxMultiHeadSelfAttention(nn.Module):
self.dim = self.config.dim
self.dropout = nn.Dropout(rate=self.config.attention_dropout)
assert self.dim % self.n_heads == 0, f"Hidden size {self.dim} not dividable by number of heads {self.n_heads}"
if not (self.dim % self.n_heads == 0):
raise ValueError(f"Hidden size {self.dim} not dividable by number of heads {self.n_heads}")
self.q_lin = nn.Dense(
self.dim,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment