torch.save(model.state_dict(), "checkpoint.pt")