mrq/torchscale
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update Bert MoE

Browse Source
This commit is contained in:
shumingma 2023-03-07 21:21:48 -08:00
parent c397ebb013
commit 0a07df1e5b
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
examples/fairseq/criterions/masked_lm_moe.py
View File

@ -29,14 +29,14 @@ class MaskedLMMoECrossEntropyCriterion(MoECriterion):
)
lprobs = model.get_normalized_probs(net_output, log_probs=True)
lprobs = lprobs.view(-1, lprobs.size(-1))
target = model.get_targets(sample, net_output).view(-1)
target = model.get_targets(sample, net_output)
if masked_tokens is not None:
targets = targets[masked_tokens]
target = target[masked_tokens]
nll_loss = F.nll_loss(
lprobs,
target,
target.view(-1),
ignore_index=self.padding_idx,
reduction="sum" if reduce else "none",
)