Muon Optimizer Update with Newton-Schulz Preconditioning - Coding problem | OneNoughtOne