Update modeling_super_linear.py

modeling_super_linear.py

@@ -357,6 +357,7 @@ class SparseNoisyMoE(nn.Module):
         
         if get_prob:
             expert_probs = F.softmax(self.gate_outputs, dim=1)
+            print(expert_probs.shape)
             return output, load_balancing_loss, expert_probs
         
         return output, load_balancing_loss
@@ -491,7 +492,7 @@ class superLinear(nn.Module):
         return cycle
 
 
-    def forward(self, x_enc, x_mark_enc=None, x_dec=None, x_mark_dec=None, mask=None, freq=[None], get_prob=False):
+    def forward(self, x_enc, x_mark_enc=None, x_dec=None, x_mark_dec=None, mask=None, freq=[None], get_prob=True):
         if len(x_enc.shape) > 2:
             x = x_enc.permute(0, 2, 1)
             B, V, L = x.shape