def METHOD_NAME():
    for m in [gpt2_medium]:
        model = m().cuda()
        model.train()
        data, mask = gen_gpt_data(GPT_BATCH_SIZE, 1024, 50257, device='meta')
        graph = ColoTracer().trace(model, meta_args={'input_ids': data, 'attention_mask': mask})
        gm = torch.fx.GraphModule(model, graph)
        interp = MetaInfoProp(gm)
        interp.propagate(MetaTensor(data, fake_device='cuda:0'), MetaTensor(mask, fake_device='cuda:0'))
        model.cpu()
        fwd_flop, bwd_flop = extract_forward_flops(gm)
        concrete_forward_mem, concrete_param_mem = run_gpt_forward(gm)
        meta_forward_mem, meta_param_mem = extract_forward_mem(gm)
        print(
            f'|{m.__name__}|{meta_forward_mem:.3f} MB|{meta_param_mem:.3f} MB|{concrete_forward_mem:.3f} MB|{concrete_param_mem:.3f} MB|fwd_flop={fwd_flop / 1e9:.3f}GFLOPs|bwd_flop={bwd_flop / 1e9:.3f}GFLOPs|'
        )
        del model, gm