fix tf function retracing for genomattn `apply`

[grahamgower]

predict_on_batch outputs the warning below. Possibly need to change the generator to produce a fixed size per batch.

5 out of the last 27 calls to <function Model.make_predict_function..predict_function at 0x7f012446a050> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.