From e6a0c99cca2a9191d33afc182360fcb8d9c851e2 Mon Sep 17 00:00:00 2001 From: lihongyang1990 Date: Wed, 24 Jun 2026 20:52:39 +0800 Subject: [PATCH] fix plugin te_general_grouped test bug --- transformer_engine/plugin/tests/test_te_general_grouped.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/plugin/tests/test_te_general_grouped.py b/transformer_engine/plugin/tests/test_te_general_grouped.py index 1bc815cc8b..defc8b4220 100644 --- 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