fix __all__ in pytorch/permutation.py

Signed-off-by: Hongxiao Bai <[email protected]>
NVIDIA · Dec 13, 2024 · ca94d72 · ca94d72
1 parent a8f1daa
commit ca94d72
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py
@@ -16,6 +16,7 @@
 __all__ = [
     "moe_permute",
     "moe_unpermute",
+    "moe_sort_chunks_by_indices",
 ]