Skip to content

GitLab

Unverified Commit 5f8e3999 authored by Fragile-azalea's avatar Fragile-azalea Committed by GitHub
Browse files

fix GshardGate top1_idx

parent f2e0d1d5
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 2 deletions
fmoe/gates/gshard_gate.py
View file @ 5f8e3999
...@@ -21,8 +21,7 @@ class GShardGate(NaiveGate): ...@@ -21,8 +21,7 @@ class GShardGate(NaiveGate):
topk_idx, topk_val, gate_score = naive_outs topk_idx, topk_val, gate_score = naive_outs
S = gate_score.shape[0] S = gate_score.shape[0]
top_k = topk_idx.shape[0] // gate_score.shape[0] top1_idx = topk_idx.view((-1, self.top_k))[:, 0]
top1_idx = topk_idx.view((-1, top_k))[:, 0]
c_e = torch.scatter_add( c_e = torch.scatter_add(
torch.zeros(self.tot_expert, device=top1_idx.device), torch.zeros(self.tot_expert, device=top1_idx.device),
0, 0,
......
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