How to properly add new submodules during training

[OlfwayAdbayIgbay]

I am currently running a module which grows during training.
Every X epochs, it grows by another submodule, and all the previous parameters are detached.

The challenges with this were as follows:

Reloading the optimizer for the new parameters:

    def on_train_epoch_end(self):
        if (self.current_epoch % self.encoder_subepoch_length) == 0:
            self._initialize_new_step()
            self.trainer.optimizers = [torch.optim.Adam(self.primary_module.parameters(), self.lr, weight_decay=self.wd)]

and expanding by another module:

    def _reset_primary_module(self):
        self.iteration_count += 1
        self.primary_module = MyModule(args).double().cuda(0)

Now I mainly have problems with this last part. The .double() and .cuda(0) calls are suboptimal and should be done via the trainer, because I also want to use things like distributed data parallel as well as experimenting with precision.

Is there a way to do this "properly"? Any methods the trainer has that could do this for me?