import torch


def load_checkpoint(path, model, optimizer=None):
    pth = torch.load(path)

    model.load_state_dict(pth["state_dict"])
    if optimizer:
        optimizer.load_state_dict(pth["optimizer"])

    print("Checkpoint {} successfully loaded".format(path))

    return pth["epoch"], pth["total_iter"]


def save_checkpoint(state, path):
    torch.save(state, path)