Remove unnecessary autograph=True.

PiperOrigin-RevId: 587350086
google-deepmind · Dec 2, 2023 · e939af5 · e939af5
1 parent c3fce27
commit e939af5
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)