Tensorboard and Trainer pattern not working together

I’m using the Keras Trainer pattern as illustrated here. The issue when using this pattern is that when you use Tensorboard only the top level weights are being recorded.

The reason for this is that Tensorboard is recording the weights for the all the layers in self.model.layers. But this equal to [<Sequential name=sequential, built=True>] and the weghts for that Sequential object is

I tried several things:

  1. Passing a CallBackList to the Tensorflow Trainer when calling fit passing model_a instead of trainer_a, but this fails because model_a has no optimizer
  2. I tried to overwrite the layers method in thr Trainer object to have recursive=True but the weights were still not showing in TensorBoard suggesting that something else is going on

I’m open to any suggestions here. I also opened an issue on the keras repo: Tensorboard not working with Trainer Pattern · Issue #20809 · keras-team/keras · GitHub