Use callback information for training neural network in tensorflow 2

Question

I am currently trying to extend the code suggested in this question: How to get layer weight while training?

By using the code below:

class CustomCallb