feat(//py): register trtorch with torch op library to support

deserializing modules in python Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · May 31, 2020 · 736e914 · 736e914
1 parent eb636ec
commit 736e914
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/py/trtorch/__init__.py b/py/trtorch/__init__.py
@@ -11,3 +11,7 @@
 from trtorch._compiler import *
 from trtorch._types import *
 from trtorch import logging
+
+def _register_with_torch():
+    trtorch_dir = os.path.dirname(__file__)
+    torch.ops.load_library(trtorch_dir + '/lib/trtorch.so')