Updating optimizer during training from callback with DDP

[fColangelo]

Hi,

I am implementing a callback for progressive resizing, basically updating the resolution of images in the datamodule and the optimizers after certain pre-defined epoch thresholds. Currently I am using the on_train_epoch_start method to do this.
However, I am not really sure how the optimizer (and scheduler) should be updated in order to not break stuff like DDP etc.
I found this issue that suggested to simply update trainer and scheduler manually like so:

   def on_train_epoch_start(self, trainer, pl_module):
       if trainer.current_epoch == self.splits[idx].ep_start:
            # Reset optimizer to adapt to the new resolution
            [...]
            trainer.optimizers    = [new_optimizers]
            trainer.lr_schedulers = trainer.configure_schedulers([new_schedulers])
            trainer.optimizer_frequencies = [] # or optimizers frequencies if you have any

Is this compatible with DDP? Is there a recommended way to do this?
Thanks!