Update modeling_super_linear.py

modeling_super_linear.py

@@ -354,6 +354,9 @@ class SparseNoisyMoE(nn.Module):
         output = torch.sum(self.topk_gates.unsqueeze(2) * sparse_expert_outputs, dim=1)
 
         load_balancing_loss = self.calculate_load_balancing_loss(self.gate_outputs, batch_size)
+
+        expert_probs = F.softmax(self.gate_outputs, dim=1)
+        print(expert_probs.shape)
         
         if get_prob:
             expert_probs = F.softmax(self.gate_outputs, dim=1)
@@ -509,8 +512,6 @@ class superLinear(nn.Module):
         else:
             out, self.moe_loss = self.moe(x)
 
-        print(out.shape)
-
 
         if self.auto_regressive and self.max_horizon < self.inf_pred_len:
             #print("bitch")