diff --git a/Project_Model/Libs/Transformer/Classes/__init__.py b/Project_Model/Libs/Transformer/Classes/__init__.py
index 837ea82..b4507f1 100644
--- a/Project_Model/Libs/Transformer/Classes/__init__.py
+++ b/Project_Model/Libs/Transformer/Classes/__init__.py
@@ -1,15 +1,16 @@
 from .Decoder import Decoder
 from .Encoder import Encoder
 from .FeedForwardNetwork import FeedForwardNetwork
-from .MultiHeadAttention import MultiheadAttention
+# from .MultiHeadAttention import MultiheadAttention
 from .TorchMultiHeadAttention import TorchMultiHeadAttention
 from .SpannedMasker import SpannedMasker
+from .DeToken import DeToken
 
 __all__ = [
     "Decoder",
     "Encoder",
     "FeedForwardNetwork",
-    "MultiheadAttention",
     "TorchMultiHeadAttention",
-    "SpannedMasker"
+    "SpannedMasker",
+    "DeToken"
 ]
\ No newline at end of file