fix

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

tests/unit_tests/_peft/test_lora_moe.py

@@ -16,6 +16,10 @@
 import torch
 import torch.nn as nn
 from unittest.mock import MagicMock, patch
+try:
+    import grouped_gemm
+except ImportError:
+    grouped_gemm = None
 
 from nemo_automodel.components.moe.layers import GroupedExperts, GroupedExpertsDeepEP, MoEConfig
 from nemo_automodel.components._peft.lora_moe import GroupedExpertsLoRA, GroupedExpertsDeepEPLoRA
@@ -156,15 +160,6 @@ def test_grouped_experts_deepep_lora_init(moe_config, device):
     assert lora_experts.lora_gate_and_up_B.requires_grad
 
 
-try:
-    import grouped_gemm
-except ImportError:
-    grouped_gemm = None
-
-
-
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_patch_moe_module(moe_config, device):
     """Test that patch_moe_module correctly wraps the original experts with the appropriate LoRA class."""