Skip to content

GitLab

Unverified Commit 4edeccd9 authored by Jiezhong Qiu's avatar Jiezhong Qiu Committed by GitHub
Browse files

Merge pull request #92 from laekov/naive-gate-fix 

Fix naive gate output score
parents 9dc9c1fa a6ca7f62
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
fmoe/gates/naive_gate.py
View file @ 4edeccd9
...@@ -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