    def set_model(self, model: nn.Module) -> None:
        """
        Set the model to callback.

        Args:
            model (`torch.nn.Module`):
                The model which callback can operate.
        """
        self.model = model
Set the model to callback.

Args:
model (`torch.nn.Module`):
The model which callback can operate.
towhee/trainer/callback.py