diff --git a/src/pystencils_autodiff/backends/_tensorflow.py b/src/pystencils_autodiff/backends/_tensorflow.py
index e0e6f0b98ece00ed00b29a2ecdf5859bb5094b4d..b24343781d40d3845b3b1a99c5b68b46e8d06040 100644
--- a/src/pystencils_autodiff/backends/_tensorflow.py
+++ b/src/pystencils_autodiff/backends/_tensorflow.py
@@ -62,12 +62,12 @@ def tensorflowop_from_autodiffop(autodiffop: pystencils_autodiff.AutoDiffOp,
         return [rtn_dict[o.name] for o in autodiffop._backward_output_fields]
 
     def backward(op, *grad):
-        return tf.py_func(helper_backward,
-                          [*op.inputs,
+        return py_func(helper_backward,
+                       [*op.inputs,
                            *grad],
-                          [f.dtype.numpy_dtype for f in autodiffop._backward_output_fields],
-                          name=autodiffop.op_name + '_backward',
-                          stateful=False)
+                       [f.dtype.numpy_dtype for f in autodiffop._backward_output_fields],
+                       name=autodiffop.op_name + '_backward',
+                       stateful=False)
 
     output_tensors = _py_func(helper_forward,
                               [inputfield_tensor_dict[f]