diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py
index d77f41c05a..4a3a811a76 100644
--- 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",
 ]