diff --git a/docs/conf.py b/docs/conf.py
index 5bc0c3b62b9ad1d200e47f8858cbbc1fec323226..5a2a0cfe56bfddb2edd791269c5cc16437fb3c94 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -141,7 +141,7 @@ pygments_style = 'sphinx'
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 html_theme = 'alabaster'
-# html_theme = 'sphinx_rtd_theme'  # official readthedocs theme used also by Martin
+# html_theme = 'sphinx_rtd_theme' #official readthedocs theme used also by Martin
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
diff --git a/src/pystencils_autodiff/backends/_torch_native.py b/src/pystencils_autodiff/backends/_torch_native.py
index 1fab51c5e80c1934fa4ba9b98137b5a5c1b99bf6..1276e734f3dfca708714614dd4c8ff26cd5fcccf 100644
--- a/src/pystencils_autodiff/backends/_torch_native.py
+++ b/src/pystencils_autodiff/backends/_torch_native.py
@@ -122,6 +122,7 @@ def generate_torch(destination_folder,
     from torch.utils.cpp_extension import load
     compiled_operation = load(operation_string, required_files, verbose=True,
                               extra_cuda_cflags=[] if is_cuda else [])
+    compiled_operation.code = output
     return compiled_operation
 
 
@@ -163,6 +164,7 @@ def create_autograd_function(autodiff_obj, inputfield_to_tensor_dict, forward_lo
         cls.saved = None
         cls.forward = forward
         cls.backward = backward
+        cls.code = compiled_operation.code
         return cls()
     else:
         op = pystencils_autodiff.backends._pytorch.create_autograd_function(autodiff_obj,