def METHOD_NAME(self, baseline_mp2, inputs, class_tmpdir):
    args_defaults = {
        'num_layers': 2,
        'hidden_size': 128,
        'num_attention_heads': 8,
        'max_position_embeddings': 128,
    }
    world_size = os.environ["WORLD_SIZE"]
    model = get_gpt2_model(args_defaults, mp_size=world_size)
    model = get_deepspeed_model(model)
    model.eval()
    with torch.no_grad():
        model.load_checkpoint(class_tmpdir, load_optimizer_states=False, load_lr_scheduler_states=False)
        device_name = get_accelerator().device_name()
        METHOD_NAME = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
        if dist.get_rank() == 0:
            load_path = os.path.join(class_tmpdir, "output.pt")
            baseline = torch.load(load_path)
            METHOD_NAME = METHOD_NAME.cpu()
            assert torch.allclose(
                baseline, METHOD_NAME,
                atol=1e-03), f"Baseline output {baseline} is not equal to save-then-load output {METHOD_NAME}"