MOE debug

.idea/workspace.xml

@@ -5,7 +5,8 @@
   </component>
   <component name="ChangeListManager">
     <list default="true" id="23565123-73ab-4f40-a9ef-1086e0c9e1ec" name="Changes" comment="">
-      <change beforePath="$PROJECT_DIR$/src/lcnn_model.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/lcnn_model.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/src/moe_model.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/moe_model.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -131,7 +132,7 @@
       <workItem from="1742889011586" duration="41000" />
       <workItem from="1742890414523" duration="4382000" />
       <workItem from="1742941362366" duration="1283000" />
-      <workItem from="1742970904084" duration="1716000" />
+      <workItem from="1742970904084" duration="2326000" />
     </task>
     <servers />
   </component>

src/moe_model.py

@@ -7,6 +7,8 @@ class UltimateMOE(nn.Module):
     def __init__(self, experts):
         super(UltimateMOE, self).__init__()
 
+        self.threshold = 0.3
+
         self.experts = nn.ModuleList(experts)
         num_experts = len(experts)
 
@@ -28,6 +30,7 @@ class UltimateMOE(nn.Module):
             nn.Linear(64, num_experts),
         )
 
+        self.softmax = nn.Softmax(dim=1)
 
     def forward(self, x):
 
@@ -49,4 +52,6 @@ class UltimateMOE(nn.Module):
         weighted_logits = torch.stack(outputs, dim=-1)
         weighted_logits = torch.einsum('bn,bcn->bc', gating_weights, weighted_logits)
 
-        return weighted_logits, gating_weights
+        score = self.softmax(weighted_logits)
+
+        return score