@staticmethod
def forward(ctx, inputs, weight, adj_matrix, am_partitions, activation_func, btype):
    ctx.save_for_backward(inputs, weight, adj_matrix)
    ctx.am_partitions = am_partitions
    ctx.activation_func = activation_func
    ctx.btype = btype
    z = broad_func(adj_matrix.size(0), am_partitions, inputs, btype=btype)
    torch.cuda.synchronize()
    g_timer.start(f'gcn_mm_ep{cur_epoch}')
    z = torch.mm(z, weight)
    torch.cuda.synchronize()
    g_timer.stop(f'gcn_mm_ep{cur_epoch}')
    z.requires_grad = True
    ctx.z = z
    return activation_func(z)
@staticmethod
def forward(ctx, inputs, weight, adj_matrix, am_partitions, activation_func, btype):
    ctx.save_for_backward(inputs, weight, adj_matrix)
    ctx.am_partitions = am_partitions
    ctx.activation_func = activation_func
    ctx.btype = btype
    z = broad_func(adj_matrix.size(0), am_partitions, inputs, btype=btype)
    torch.cuda.synchronize()
    g_timer.start(f'gcn_mm_ep{cur_epoch}')
    z = torch.mm(z, weight)
    torch.cuda.synchronize()
    g_timer.stop(f'gcn_mm_ep{cur_epoch}')
    z.requires_grad = True
    ctx.z = z
    return activation_func(z)