Skip to content

Commit 9aa7e20

Browse files
[fix] fix plugin te_general_grouped test bug (#81)
## Summary Replace the grouped GEMM plugin test import from `transformer_engine_torch_nv` to `transformer_engine_torch`. ## Why The test should load the standard Transformer Engine torch extension module. Importing `transformer_engine_torch_nv` can fail because that is not the expected module name in this project. ## Impact This fixes startup/import failure for `transformer_engine/plugin/tests/test_te_general_grouped.py` before the grouped GEMM checks run. ## Validation - `git diff --check origin/main...HEAD` - `python3 -c 'import ast, pathlib; path = pathlib.Path("transformer_engine/plugin/tests/test_te_general_grouped.py"); ast.parse(path.read_text())'` Full test execution was not run locally because it requires the project runtime/CUDA environment.
1 parent b7f65d1 commit 9aa7e20

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

transformer_engine/plugin/tests/test_te_general_grouped.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def test_grouped_gemm_equivalence(self, grad, has_bias, has_pre_gelu, single_out
2222
"\n test te_general_grouped_gemm"
2323
f" grad:{grad} has_bias:{has_bias},has_pre_gelu:{has_pre_gelu},single_output:{single_output}"
2424
)
25-
import transformer_engine_torch_nv as tex
25+
import transformer_engine_torch as tex
2626

2727
num_gemms = 2
2828
m, k, n = 128, 32, 64

0 commit comments

Comments
 (0)