Update modeling_super_linear.py

modeling_super_linear.py

@@ -356,6 +356,17 @@ 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)
+        import matplotlib.pyplot as plt
+
+        plt.figure(figsize=(10, 6))
+        plt.bar(range(len(expert_probs)), expert_probs.detach().cpu().numpy())
+        plt.xlabel('Expert Index')
+        plt.ylabel('Probability')
+        plt.title('Expert Selection Probabilities')
+        plt.xticks(range(len(expert_probs)), [f'Expert {i}' for i in range(len(expert_probs))])
+        plt.tight_layout()
+        plt.savefig('expert_probabilities.png')
+        plt.close()
         print(expert_probs.shape)
         
         if get_prob: