[Bug]: Colab x Weights & Biases Keras Callbacks not working

[Esmedd]

Tell us what happened?

Example - Use_WandbEvalCallback_in_your_Keras_workflow.ipynb

Error message at cell 2

ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-2-3d5089b1fc8b>](https://localhost:8080/#) in <cell line: 10>()
      8 # Weights and Biases related imports
      9 import wandb
---> 10 from wandb.keras import WandbMetricsLogger
     11 from wandb.keras import WandbModelCheckpoint
     12 from wandb.keras import WandbEvalCallback

ModuleNotFoundError: No module named 'wandb.keras'

This has also happened in other, personal colabs despite restarting the kernel after install, running uninstall/install, upgrading etc. Also have tried pip3 to no avail.

[tcapelle]

cc @soumik12345

[tcapelle]

If you revert to wandb<0.17 it works, I will check where are those imports now and update the colab. Or:

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras import models
import tensorflow_datasets as tfds

# Weights and Biases related imports
import wandb
# from wandb.keras import WandbMetricsLogger
# from wandb.keras import WandbModelCheckpoint
# from wandb.keras import WandbEvalCallback
from wandb.integration.keras import WandbMetricsLogger
from wandb.integration.keras import WandbModelCheckpoint
from wandb.integration.keras import WandbEvalCallback

[Esmedd]

Great thanks, it works now :)