tfa.optimizers.MultiOptimizer does not work on TPU
artitechture opened this issue · 0 comments
artitechture commented
System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): google colab
- TensorFlow version and how it was installed (source or binary): 2.12.0
- TensorFlow-Addons version and how it was installed (source or binary): 0.20.0
- Python version: 3.9
- Is GPU used? (yes/no): TPU used
Describe the bug
tfa.optimizers.MultiOptimizer raise InvalidArgumentError on TPU. The following code runs without a error on GPU.
Code to reproduce the issue
Other info / logs
---------------------------------------------------------------------------
InvalidArgumentError Traceback (most recent call last)
<ipython-input-8-521e74c800a4> in <cell line: 1>()
----> 1 model.fit(x_train, y_train, validation_data=(x_test, y_test),
2 epochs=10, batch_size=16*strategy.num_replicas_in_sync)
1 frames
/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
68 # To get the full stack trace, call:
69 # `tf.debugging.disable_traceback_filtering()`
---> 70 raise e.with_traceback(filtered_tb) from None
71 finally:
72 del filtered_tb
/usr/local/lib/python3.9/dist-packages/tensorflow/python/framework/ops.py in _numpy(self)
1126 return self._numpy_internal()
1127 except core._NotOkStatusException as e: # pylint: disable=protected-access
-> 1128 raise core._status_to_exception(e) from None # pylint: disable=protected-access
1129
1130 @property
InvalidArgumentError: Missing edge input number: 7