update mha dim

This commit is contained in:
rasbt 2024-05-27 07:46:29 -05:00
parent 68123a5433
commit c0f564ee87

View File

@ -239,7 +239,7 @@
" self.head_dim = d_out // num_heads\n",
"\n",
" self.qkv = nn.Linear(d_in, 3 * d_out, bias=qkv_bias)\n",
" self.proj = nn.Linear(d_in, d_out)\n",
" self.proj = nn.Linear(d_out, d_out)\n",
" self.dropout = nn.Dropout(dropout)\n",
"\n",
" self.register_buffer(\n",