diff --git a/src/pyronn_torch/codegen.py b/src/pyronn_torch/codegen.py
index 755ac01bf2204fa1c35e68d83b919671977a27c6..f4a3cfa841957edb0e8cb92a527c5db2756180eb 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():