Using WandB with Keras without

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 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 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 and this is not the case. Can you help please!
Thanks a lot for your help.