[fbsync] Enabling torchtext extension work seamlessly between OSS and…

… on fbcode (#1382)
pytorch · Aug 27, 2021 · aa12e9a · aa12e9a
1 parent 325e235
commit aa12e9a
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 20 deletions.
diff --git a/torchtext/__init__.py b/torchtext/__init__.py
@@ -5,6 +5,7 @@
 from . import vocab
 from . import experimental
 from . import legacy
+from ._extension import _init_extension
 
 
 try:
@@ -21,26 +22,6 @@
            'legacy']
 
 
-def _init_extension():
-    import os
-    import importlib
-    import torch
-
-    # load the custom_op_library and register the custom ops
-    lib_dir = os.path.dirname(__file__)
-    loader_details = (
-        importlib.machinery.ExtensionFileLoader,
-        importlib.machinery.EXTENSION_SUFFIXES
-    )
-
-    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
-    ext_specs = extfinder.find_spec("_torchtext")
-    if ext_specs is None:
-        raise ImportError("torchtext C++ Extension is not found.")
-    torch.ops.load_library(ext_specs.origin)
-    torch.classes.load_library(ext_specs.origin)
-
-
 _init_extension()
 
 

diff --git a/torchtext/_extension.py b/torchtext/_extension.py
@@ -0,0 +1,18 @@
+def _init_extension():
+    import os
+    import importlib
+    import torch
+
+    # load the custom_op_library and register the custom ops
+    lib_dir = os.path.dirname(__file__)
+    loader_details = (
+        importlib.machinery.ExtensionFileLoader,
+        importlib.machinery.EXTENSION_SUFFIXES
+    )
+
+    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
+    ext_specs = extfinder.find_spec("_torchtext")
+    if ext_specs is None:
+        raise ImportError("torchtext C++ Extension is not found.")
+    torch.ops.load_library(ext_specs.origin)
+    torch.classes.load_library(ext_specs.origin)