diff --git a/fastNLP/embeddings/__init__.py b/fastNLP/embeddings/__init__.py
index 4f90ac63..37881f17 100644
--- a/fastNLP/embeddings/__init__.py
+++ b/fastNLP/embeddings/__init__.py
@@ -7,6 +7,7 @@ torch.FloatTensor。所有的embedding都可以使用 `self.num_embedding` 获
 
 __all__ = [
     "Embedding",
+    "TokenEmbedding",
     "StaticEmbedding",
     "ElmoEmbedding",
     "BertEmbedding",
@@ -14,14 +15,14 @@ __all__ = [
     "StackEmbedding",
     "LSTMCharEmbedding",
     "CNNCharEmbedding",
-    "get_embeddings"
+    "get_embeddings",
 ]
 
 
-from .embedding import Embedding
+from .embedding import Embedding, TokenEmbedding
 from .static_embedding import StaticEmbedding
 from .elmo_embedding import ElmoEmbedding
 from .bert_embedding import BertEmbedding, BertWordPieceEncoder
 from .char_embedding import CNNCharEmbedding, LSTMCharEmbedding
 from .stack_embedding import StackEmbedding
-from .utils import get_embeddings
\ No newline at end of file
+from .utils import get_embeddings