I have a custom layer which uses tf.py_function with some python code. Since this isn't supported by Cloud TPU, is it possible to place the computation for this layer on a CPU or GPU device? Or is the only solution to rewrite the function?
This use case is not currently supported.