add thres

.idea/workspace.xml

@@ -6,7 +6,6 @@
   <component name="ChangeListManager">
     <list default="true" id="23565123-73ab-4f40-a9ef-1086e0c9e1ec" name="Changes" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/script.py" beforeDir="false" afterPath="$PROJECT_DIR$/script.py" 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" />
@@ -137,7 +136,7 @@
       <workItem from="1743062628099" duration="35000" />
       <workItem from="1743063082652" duration="7000" />
       <workItem from="1743092790258" duration="1395000" />
-      <workItem from="1743151940209" duration="996000" />
+      <workItem from="1743151940209" duration="1035000" />
     </task>
     <servers />
   </component>

src/moe_model.py

@@ -62,6 +62,8 @@ class MOE_attention(nn.Module):
     def __init__(self, experts, device, input_dim=128, freezing=False):
         super(MOE_attention, self).__init__()
 
+        self.threshold = 0.4
+
         self.device = device
         self.experts = nn.ModuleList(experts)
         self.num_experts = len(experts)