From d1c464ac3834499b3e005cd0be9d5e861efbeb00 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 24 Feb 2020 17:16:10 +0100
Subject: [PATCH] Return JIT compiled extension
---
src/pyronn_torch/codegen.py | 6 ++----
1 file changed, 2 insertions(+), 4 deletions(-)
diff --git a/src/pyronn_torch/codegen.py b/src/pyronn_torch/codegen.py
index 755ac01..f4a3cfa 100644
--- a/src/pyronn_torch/codegen.py
+++ b/src/pyronn_torch/codegen.py
@@ -178,14 +178,12 @@ def generate_shared_object(output_folder=None, source_files=None, show_code=Fals
extension = module.compile(extra_source_files=cuda_sources, extra_cuda_flags=['-arch=sm_35'], with_cuda=True)
- for v in extension.__dict__.values():
- if hasattr(v, '__doc__'):
- print(v.__doc__)
-
shared_object_file = module.compiled_file.replace('.cpp', '.so')
copyfile(shared_object_file, join(output_folder, 'pyronn_torch.so'))
copyfile(module.compiled_file, join(output_folder, 'pyronn_torch.cpp'))
+ return extension
+
def main():
--
GitLab