muon imports in test v2

Author: SalmanMohammadi

tests/core/test_builders.py

@@ -475,6 +475,7 @@ def test_custom_optimizer_cls_and_kwargs(
             assert trainer.optimizer_cls_and_kwargs is not None
 
             from axolotl.contribs.mit.muon import MuonOptimizerFactory
+            from axolotl.contribs.mit.muon.muon import Muon
 
             optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
             assert optimizer_cls is MuonOptimizerFactory
@@ -554,6 +555,7 @@ def test_custom_optimizer_cls_and_kwargs(
         assert trainer.optimizer_cls_and_kwargs is not None
 
         from axolotl.contribs.mit.muon import MuonOptimizerFactory
+        from axolotl.contribs.mit.muon.muon import Muon
 
         optimizer_cls, optimizer_kwargs = trainer.optimizer_cls_and_kwargs
         assert optimizer_cls is MuonOptimizerFactory