Error with Tensorflow 2.6.0

Hi,
I had a working code in Google Colab with Tensorflow 2.5.0. Colab has recently updated the Tensorflow version to 2.6.0 and the code is now giving the following error:

TypeError: ray.cloudpickle.dumps(<class 'ray.tune.function_runner.wrap_function.<locals>.ImplicitFunc'>) failed.
To check which non-serializable variables are captured in scope, re-run the ray script with 'RAY_PICKLE_VERBOSE_DEBUG=1'. Other options: 
-Try reproducing the issue by calling `pickle.dumps(trainable)`. 
-If the error is typing-related, try removing the type annotations and try again.

Traceback points the problem at the partial funtion call:
analysis = tune.run(partial (train_hpo_func, ...)

Appreciate it if someone could help fix this. Thanks.

In case some others are also facing this, the problem is resolved when keras.api._v2.keras is imported in the module containing tune.run.