What's wrong @tf.py_function

I tried code the sample in colab:

def py_log_huber(x, m):
print(‘Running with eager execution.’)
if tf.abs(x) <= m:
return x2
return m
2 * (1 - 2 * tf.math.log(m) + tf.math.log(x**2))

TypeError Traceback (most recent call last)
in <cell line: 1>()
----> 1 @tf.py_function(Tout=tf.float32)
2 def py_log_huber(x, m):
3 print(‘Running with eager execution.’)
4 if tf.abs(x) <= m:
5 return x**2

/usr/local/lib/python3.10/dist-packages/tensorflow/python/util/dispatch.py in op_dispatch_handler(*args, **kwargs)
1168 if iterable_params is not None:
1169 args, kwargs = replace_iterable_params(args, kwargs, iterable_params)
→ 1170 result = api_dispatcher.Dispatch(args, kwargs)
1171 if result is not NotImplemented:
1172 return result

TypeError: Missing required positional argument

Hi @hsiaokai0309

Welcome to The TensorFlow Forum!

Could you please try executing the code again in Colab because the above tf.py_function code is working as expected in Google Colab. You have already mentioned the required argument in the decorator so there is no missing argument.

Please have a look at the below screenshot.

1 Like