Skip to content
Snippets Groups Projects
Commit 3b6ec0fd authored by Stephan Seitz's avatar Stephan Seitz
Browse files

Improve docstring of TensorflowModule

parent 43cb2402
No related branches found
No related tags found
No related merge requests found
......@@ -141,8 +141,10 @@ class TensorflowModule(TorchModule):
def __init__(self, module_name, kernel_asts):
"""Create a C++ module with forward and optional backward_kernels
:param forward_kernel_ast: one or more kernel ASTs (can have any C dialect)
:param backward_kernel_ast:
Args:
module_name (str): Module name
kernel_asts (pystencils.kernel_wrappers.KernelWrapper):
ASTs as generated by `:func:pystencils.create_kernel`
"""
self.compiled_file = None
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment