Enforce autograph=False on gradient function.

PiperOrigin-RevId: 587754205
google-deepmind · Dec 4, 2023 · 268b7dc · 268b7dc
1 parent c3fce27
commit 268b7dc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tf2jax/_src/tf2jax.py b/tf2jax/_src/tf2jax.py
@@ -1405,7 +1405,7 @@ def _convert_gradient_function(
     library[grad_fn_name] = None
     return
 
-  @tf.function
+  @tf.function(autograph=False)
   def tf_grad_fn(*grad_args, **grad_kwargs):
     fn = tf_ops.gradient_registry.lookup(grad_fn_name)
     return fn(None, *grad_args, **grad_kwargs)