fix switch_optimizer callback

pina/callback/optimizer_callback.py

@@ -21,26 +21,30 @@ class SwitchOptimizer(Callback):
             single :class:`torch.optim.Optimizer` instance or a list of them
             for multiple model solver.
         :type new_optimizers: pina.optim.TorchOptimizer | list
-        :param epoch_switch: The epoch at which the optimizer switch occurs.
-        :type epoch_switch: int
+        :param int epoch_switch: The epoch at which the optimizer switch occurs.
 
         Example:
-            >>> switch_callback = SwitchOptimizer(new_optimizers=optimizer,
-            >>>                                  epoch_switch=10)
+            >>> optimizer = TorchOptimizer(torch.optim.Adam, lr=0.01)
+            >>> switch_callback = SwitchOptimizer(
+            >>>     new_optimizers=optimizer, epoch_switch=10
+            >>> )
         """
         super().__init__()
 
+        # Check if epoch_switch is greater than 1
         if epoch_switch < 1:
             raise ValueError("epoch_switch must be greater than one.")
 
+        # If new_optimizers is not a list, convert it to a list
         if not isinstance(new_optimizers, list):
             new_optimizers = [new_optimizers]
 
-        # check type consistency
+        # Check consistency
+        check_consistency(epoch_switch, int)
         for optimizer in new_optimizers:
             check_consistency(optimizer, TorchOptimizer)
-        check_consistency(epoch_switch, int)
-        # save new optimizers
+
+        # Store the new optimizers and epoch switch
         self._new_optimizers = new_optimizers
         self._epoch_switch = epoch_switch
 
@@ -48,18 +52,21 @@ class SwitchOptimizer(Callback):
         """
         Switch the optimizer at the start of the specified training epoch.
 
-        :param trainer: The trainer object managing the training process.
-        :type trainer: pytorch_lightning.Trainer
+        :param lightning.pytorch.Trainer trainer: The trainer object managing
+            the training process.
         :param _: Placeholder argument (not used).
-
-        :return: None
-        :rtype: None
         """
+        # Check if the current epoch matches the switch epoch
         if trainer.current_epoch == self._epoch_switch:
             optims = []
 
+            # Hook the new optimizers to the model parameters
             for idx, optim in enumerate(self._new_optimizers):
                 optim.hook(trainer.solver._pina_models[idx].parameters())
                 optims.append(optim)
 
+            # Update the solver's optimizers
             trainer.solver._pina_optimizers = optims
+
+            # Update the trainer's strategy optimizers
+            trainer.strategy.optimizers = [o.instance for o in optims]