Open artitechture opened 1 year ago
System information
Describe the bug
tfa.optimizers.MultiOptimizer raise InvalidArgumentError on TPU. The following code runs without a error on GPU.
Code to reproduce the issue
https://colab.research.google.com/gist/artitechture/12ddf883cca36290f9580bbbb804df27/tpu_multioptimizer_cifar10.ipynb
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
System information
Describe the bug
tfa.optimizers.MultiOptimizer raise InvalidArgumentError on TPU. The following code runs without a error on GPU.
Code to reproduce the issue
https://colab.research.google.com/gist/artitechture/12ddf883cca36290f9580bbbb804df27/tpu_multioptimizer_cifar10.ipynb
Other info / logs