Update models.py

mlcommons · Jan 16, 2024 · 2458567 · 2458567
1 parent 1363e15
commit 2458567
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/algorithmic_efficiency/workloads/imagenet_vit/imagenet_jax/models.py b/algorithmic_efficiency/workloads/imagenet_vit/imagenet_jax/models.py
@@ -154,7 +154,7 @@ def __call__(self, x):
 
     x = nn.MultiHeadDotProductAttention(
         num_heads=self.num_heads,
-        use_bias=False,
+        use_bias=True,
         kernel_init=nn.initializers.xavier_uniform())(probe, x)
 
     y = nn.LayerNorm()(x)