fix naive gate output score

a6ca7f62 · Rick Ho · 9dc9c1fa · a6ca7f62
Commit a6ca7f62 authored Nov 29, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

fmoe/gates/naive_gate.py fmoe/gates/naive_gate.py +2 -2

No files found.
--- a/fmoe/gates/naive_gate.py
+++ b/fmoe/gates/naive_gate.py
@@ -38,5 +38,5 @@ class NaiveGate(BaseGate):
        gate_score = F.softmax(gate_top_k_val, dim=-1)

        if return_all_scores:
-            return gate_top_k_idx, gate_top_k_val, gate
-        return gate_top_k_idx, gate_top_k_val
+            return gate_top_k_idx, gate_score, gate
+        return gate_top_k_idx, gate_score