def save_model(self, model_dir=None, model=None, optimizer=None):
        """Save the model and the optim."""
        if not model_dir:
            model_dir = self.model_dir
        os.makedirs(model_dir, exist_ok=True)

        if model:
            # Take care of distributed/parallel training
            torch.save(model.state_dict(), self.checkpoint_file)
        if optimizer:
            torch.save(optimizer.state_dict(), self.optimizer_file)
Save the model and the optim.
pycorrector/deepcontext/deepcontext_model.py