fix plugin te_general_grouped test bug

lihongyang1990 · lihongyang1990 · commit e6a0c99cca2a · 2026-06-24T20:52:39.000+08:00
diff --git a/transformer_engine/plugin/tests/test_te_general_grouped.py b/transformer_engine/plugin/tests/test_te_general_grouped.py
@@ -22,7 +22,7 @@ def test_grouped_gemm_equivalence(self, grad, has_bias, has_pre_gelu, single_out
             "\n test te_general_grouped_gemm"
             f" grad:{grad} has_bias:{has_bias},has_pre_gelu:{has_pre_gelu},single_output:{single_output}"
         )
-        import transformer_engine_torch_nv as tex
+        import transformer_engine_torch as tex
 
         num_gemms = 2
         m, k, n = 128, 32, 64

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ def test_grouped_gemm_equivalence(self, grad, has_bias, has_pre_gelu, single_out`
`22`	`22`	`"\n test te_general_grouped_gemm"`
`23`	`23`	`f" grad:{grad} has_bias:{has_bias},has_pre_gelu:{has_pre_gelu},single_output:{single_output}"`
`24`	`24`	`)`
`25`		`- import transformer_engine_torch_nv as tex`
	`25`	`+ import transformer_engine_torch as tex`
`26`	`26`
`27`	`27`	`num_gemms = 2`
`28`	`28`	`m, k, n = 128, 32, 64`