Optimize diagnostics performance and robustness

[Vyomkesh13]

Summary

This PR improves the performance and robustness of the diagnostics utilities
without changing any numerical results or affecting the training or evaluation
pipeline.

Key Improvements

• Vectorized mean activation computation to remove Python loops and unnecessary
  GPU→CPU synchronizations
• Ensured diagnostics run strictly in inference mode to reduce memory overhead
• Disabled interactive plotting for faster execution in batch, CI, and headless
  environments
• Simplified square-matrix construction logic while preserving identical output
• Added lightweight input validation to prevent silent misuse

Impact

• No changes to training behavior, loss functions, or evaluation metrics
• Identical diagnostic outputs and plots
• Faster and more scalable diagnostics for large models and datasets
• Improved reliability in automated and headless workflows