Update modeling_super_linear.py

modeling_super_linear.py

@@ -355,6 +355,7 @@ class SparseNoisyMoE(nn.Module):
 
         load_balancing_loss = self.calculate_load_balancing_loss(self.gate_outputs, batch_size)
 
+        expert_probs = F.softmax(self.gate_outputs, dim=1)
         expert_probs = expert_probs[1,:]
         # Plot the expert probabilities
         import matplotlib.pyplot as plt