Set more precise shape to the attention weights and outputs

modeling_backpack_gpt2.py

@@ -101,13 +101,14 @@ class BackpackWeightNetwork(nn.Module):
     super().__init__()
     self.n_embd = embed_dim
     self.num_senses = num_senses
-    self.c_attn = nn.Linear(embed_dim, 2*embed_dim)
+    self.embed_per_sense = embed_dim // num_senses
+    self.c_attn = nn.Linear(embed_dim, 2 * num_senses * self.embed_per_sense)
     self.softmax_scale = None
 
   def forward(self, encoded):
     b, s, d = encoded.shape
     encoded = self.c_attn(encoded) # (b, s, 2*d)
-    encoded = encoded.reshape(b, s, 2, self.num_senses, d // self.num_senses) #(b, s, 2, nv, d//nv)
+    encoded = encoded.reshape(b, s, 2, self.num_senses, self.embed_per_sense) #(b, s, 2, nv, d//nv)
     batch_size, seqlen = encoded.shape[0], encoded.shape[1]
 
     # compute scores & mask