Skip to content

GitLab

"tests/vscode:/vscode.git/clone" did not exist on "03406bd96d4d01b5a32378937a176cb434642519"
Commit a6ca7f62 authored by Rick Ho's avatar Rick Ho
Browse files

fix naive gate output score

parent 9dc9c1fa
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
fmoe/gates/naive_gate.py
View file @ a6ca7f62
...@@ -38,5 +38,5 @@ class NaiveGate(BaseGate): ...@@ -38,5 +38,5 @@ class NaiveGate(BaseGate):
gate_score = F.softmax(gate_top_k_val, dim=-1) gate_score = F.softmax(gate_top_k_val, dim=-1)
if return_all_scores: if return_all_scores:
return gate_top_k_idx, gate_top_k_val, gate return gate_top_k_idx, gate_score, gate
return gate_top_k_idx, gate_top_k_val return gate_top_k_idx, gate_score
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