def METHOD_NAME(self, name=None):
    """Saves the current model"""
    state = {
        "net": self.model.state_dict(),
        "optim": self.optim.state_dict(),
        "epoch": self.epochs,
        "train_accuracy": self.train_acc,
        "test_accuracy": self.valid_acc,
    }
    if name is None:
        torch.save(state, f"{self.output_dir}/{self.epochs}.pt")
    else:
        torch.save(state, f"{self.output_dir}/{name}.pt")