def _assert_checkpoint_equal(model, output_path):
    assert os.path.exists(output_path)
    single_output = torch.load(output_path)
    state_dict = model.state_dict()
    for (orig_param, saved_model_param) in zip(state_dict.values(), single_output['state_dict'].values()):
        if (model.dtype == torch.half):
            saved_model_param = saved_model_param.half()
        assert torch.equal(orig_param.cpu(), saved_model_param)
def _assert_checkpoint_equal(model, output_path):
    assert os.path.exists(output_path)
    single_output = torch.load(output_path)
    state_dict = model.state_dict()
    for (orig_param, saved_model_param) in zip(state_dict.values(), single_output['state_dict'].values()):
        if (model.dtype == torch.half):
            saved_model_param = saved_model_param.half()
        assert torch.equal(orig_param.cpu(), saved_model_param)