From documentation prepare() is used to send model, optimizer, data loaders to each TPU core etc…
If I want to save the model I will unwrap the model first by doing unwrap_model().
Is there a function to unwrap the optimizer? I assume unwrap_model() is just a function to detach from all cores? So I can use it also like unwrap_model(optimizer)
Thanks