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

Append _cuda to tensorflow op name if using cuda

parent 3cf1080f
Branches
Tags
No related merge requests found
Pipeline #18164 canceled
...@@ -52,6 +52,8 @@ def native_tensorflowop_from_autodiffop(autodiff_obj: pystencils_autodiff.AutoDi ...@@ -52,6 +52,8 @@ def native_tensorflowop_from_autodiffop(autodiff_obj: pystencils_autodiff.AutoDi
backward_ast = autodiff_obj.backward_ast_cpu backward_ast = autodiff_obj.backward_ast_cpu
autodiff_obj.op_name = f'{autodiff_obj.op_name}_{_hash(str(autodiff_obj).encode()).hexdigest()}' autodiff_obj.op_name = f'{autodiff_obj.op_name}_{_hash(str(autodiff_obj).encode()).hexdigest()}'
if use_cuda:
autodiff_obj.op_name += '_cuda'
forward_ast.function_name = autodiff_obj.op_name + "_forward" forward_ast.function_name = autodiff_obj.op_name + "_forward"
backward_ast.function_name = autodiff_obj.op_name + "_backward" backward_ast.function_name = autodiff_obj.op_name + "_backward"
module = TensorflowModule(autodiff_obj.op_name, [forward_ast, backward_ast]) module = TensorflowModule(autodiff_obj.op_name, [forward_ast, backward_ast])
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment