fix: update w_o shape

instadeepai · Nov 29, 2024 · 6d1016f · 6d1016f
1 parent cb6ff06
commit 6d1016f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mava/networks/retention.py b/mava/networks/retention.py
@@ -242,7 +242,7 @@ def setup(self) -> None:
         self.w_o = self.param(
             "w_o",
             nn.initializers.normal(stddev=1 / self.embed_dim),
-            (self.head_size, self.embed_dim),
+            (self.embed_dim, self.embed_dim),
         )
         self.group_norm = nn.GroupNorm(num_groups=self.n_head)