def test_GSA(args, input_len):
    """Test the time and CUDA memory consumption of PAM."""
    handle = pynvml.nvmlDeviceGetHandleByIndex(1)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
    init_mem = meminfo.used / 1024**3

    GSA_Layer = GraphSelfAttention(args).to(args.device)
    optimizer = optim.Adam(GSA_Layer.parameters(), 1e-4)
    optimizer.zero_grad()
    hidden_state = torch.ones(4, input_len, args.d_model, dtype=torch.float32, device=args.device)
    fake_gt = torch.zeros(4, input_len, args.d_model, device=args.device)

    # Preload the layer
    result = GSA_Layer(hidden_state)
    loss = ((fake_gt  - result) ** 2).mean()
    loss.backward()
    optimizer.step()

    used_memory = 0
    repeat_times = 1000
    start_time = time.time()
    for i in range(repeat_times):
        result = GSA_Layer(hidden_state)
        handle = pynvml.nvmlDeviceGetHandleByIndex(1)
        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        used_memory += meminfo.used / 1024**3
        loss = ((fake_gt  - result) ** 2).mean()
        loss.backward()
        optimizer.step()

    print('GSA used time:{} s'.format(round((time.time() - start_time) / repeat_times, 4)))
    used_memory = used_memory / repeat_times
    print('GSA used average memory: {} GB'.format(round(used_memory-init_mem, 4)))
Test the time and CUDA memory consumption of PAM.
PatchTST_supervised/Formers/Pyraformer/pyraformer/graph_attention.py