    def set_optimizer(self, optimizer: Optimizer) -> None:
        """
        Set the optimizer to callback.

        Args:
            optimizer (`torch.optim.Optimizer`):
                The optimizer which callback can operate.
        """
        self.optimizer = optimizer
Set the optimizer to callback.

Args:
optimizer (`torch.optim.Optimizer`):
The optimizer which callback can operate.
towhee/trainer/callback.py