Logging state of AdamW optimizer at certain stage in training

Question

I\'m using an AdamW optimizer to train a model (Tensorflow 2.0). I need to log its state at certain steps during training, so to be able to reload the optimizer\'s state and res