Compatibility with tensorflow 2.0

wandb --version && python --version && uname

• Weights and Biases version: 0.8.16
• Python version: 3.7.4
• Operating System: Win 10

Description

I try to log my custom losses using tensorboard. It seems wandb isn’t compatible with the @tf.function decorator, and it would mean that using wandb would slow down the code considerably. Is tf.function supported ?

What I Did

@tf.function
def train_on_batch(self, obs1, obs2, act, rew, done, step):
[...]
   wandb.tensorflow.log('losses/q1_loss', q1_loss, step=step)

Error: TypeError: To be compatible with tf.contrib.eager.defun, Python functions must return zero or more Tensors; in compilation of <function create_converted_entity_factory.<locals>.create_converted_entity.<locals>.tf__log.<locals>.<lambda> at 0x000001D386C6B288>, found return value of type <class ‘wandb.history.History’>, which is not a Tensor.

[raubitsj]

@Danmou, the branch should be fixed now, there was an issue how the “pip install …#egg=wandb” command was messing with the python environment.

Please try the: pip install --upgrade git+git://github.com/wandb/client.git@fix/tensorboard-import-issue#egg=wandb command again. you might want to pip uninstall wandb first just to make sure it grabs a new version.

[fellhorn]

Still having the same issue with TF 2.1.1 and TensorBoard 2.1.1. I run

import tensorflow as tf
import wandb
wandb.init(sync_tensorboard=True)
writer = tf.summary.create_file_writer('logdir')

I had the same issue when using TF and TensorBoard from the global site packages. When I switched to a venv on the machine it worked though.