Skip to content

GitLab

Commit 7945bce6 authored by Jiezhong Qiu's avatar Jiezhong Qiu
Browse files

fix bug

parent 85306aa7
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
pytorch/mem_transformer.py
View file @ 7945bce6
...@@ -87,7 +87,7 @@ class CustomizedMoEPositionwiseFF(nn.Module): ...@@ -87,7 +87,7 @@ class CustomizedMoEPositionwiseFF(nn.Module):
# x = F.pad(x, pad=(0, 1), mode='constant', value=1.0) # x = F.pad(x, pad=(0, 1), mode='constant', value=1.0)
x = self.moe2(x, gate_idx) x = self.moe2(x, gate_idx)
x = self.dropout(x) # (BxL) x d_model x = self.dropout(x) # (BxL) x d_model
output.append(x.unsqueeze(1)) # (BxL) x 1 x d_model core_out.append(x.unsqueeze(1)) # (BxL) x 1 x d_model
core_out = torch.cat(core_out, dim=1) # (BxL) x top_k x d_model core_out = torch.cat(core_out, dim=1) # (BxL) x top_k x d_model
core_out = torch.bmm(gate_score, core_out) # (BxL) x 1 x d_model core_out = torch.bmm(gate_score, core_out) # (BxL) x 1 x d_model
......
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