import torch


def load_checkpoint(checkpoint_file, model, model_opt):
    """
    Loads a checkpoint including model state and running loss for continued training
    """
    if checkpoint_file is not None:
        checkpoint = torch.load(checkpoint_file)
        state_dict = checkpoint["state_dict"]
        start_iter = checkpoint['iter']
        running_loss = checkpoint['running_loss']
        opt_state_dict = checkpoint['optimizer']
        model_opt.load_state_dict(opt_state_dict)
        for state in model_opt.state.values():
            for key, value in state.items():
                if isinstance(value, torch.Tensor):
                    state[key] = value.cuda()
        model.load_state_dict(state_dict)
    else:
        start_iter = 1
        running_loss = 0
    return start_iter, running_loss