Using WandB with Keras without model.fit

Hi everyone,
I am facing problem with implementing WandB for tracking my RL training algorithm. I am using Keras for writing my training code. However, I do not use the model.fit function. I am instead using @tf.function to perform updates on the Q-network and the Policy-network (I am using DDPG learning algorithm). How can I initialize a run and use the tracking callbacks if I am not using the model.fit so I cannot pass the callback as an argument. Every guide/tutorial on the internet for using Keras with WandB assumes that I am using model.fit and this is not the case. Can you help please!
Thanks a lot for your help.

Hello Abdelrahman Abdelgawad,

Thank you for reaching out. After reviewing our documentation on integrating wandb with Keras, as well as the Keras documentation itself, it appears that using wandb with Keras effectively requires the use of the model.fit method. This method allows you to pass the built-in callback function, which is essential for integrating wandb into your Keras workflow.

If you need further clarification or have any other questions, feel free to ask!

Best regards,
Jason

Thank you for your response. I was able to use the same integration that is used for PyTorch since I am writing the code as a bunch of tf.functions