CODEWITHSUNDEEP
pythontensorflowtensorflow2.0
Charles Staats
Charles Staats

Reputation: 203

Missing gradient when using tf.function

I have found that if I want to use tf.gradients in tensorflow 2 rather than a gradient tape, I can do this by wrapping the code in a tf.function-decorated function. But somehow, I can't take the gradient with respect to a variable this way:

import tensorflow as tf
a = tf.Variable(initial_value=1.0, dtype=tf.float32)
b = 0.01 * a

@tf.function
def get_grads():
  return tf.gradients(b, a)[0]

print(get_grads())

I would expect to get some sort of tensor as a result, a tensor that ought to evaluate to 0.01. But instead I get None. Note that I am running this on Google Colab, so there should not be any issues with the tensorflow version or installation.

What am I doing wrong?

Upvotes: 0

Views: 372

Answers (1)

Tou You
Tou You

Reputation: 1194

The op b = 0.01 * a is out the graph created by the tf.function-decorated function.

you can use :

a = tf.Variable(initial_value=1.0, dtype=tf.float32)
@tf.function
def get_grads():
  
   b = 0.01 * a
   return tf.gradients(b, a)

print(get_grads())

Upvotes: 1

Related Questions