def outer_product2(inputs, ag):
    torch.cuda.synchronize()
    g_timer.start(f'gcn_mm_ep{cur_epoch}')
    grad_weight = torch.mm(inputs, ag)
    torch.cuda.synchronize()
    g_timer.stop(f'gcn_mm_ep{cur_epoch}')
    g_timer.start(f'gcn_allreduce_ep{cur_epoch}')
    dist.all_reduce(grad_weight, op=dist.ReduceOp.SUM, group=g_env.world_group)
    torch.cuda.synchronize()
    g_timer.stop(f'gcn_allreduce_ep{cur_epoch}')
    return grad_weight
def outer_product2(inputs, ag):
    torch.cuda.synchronize()
    g_timer.start(f'gcn_mm_ep{cur_epoch}')
    grad_weight = torch.mm(inputs, ag)
    torch.cuda.synchronize()
    g_timer.stop(f'gcn_mm_ep{cur_epoch}')
    g_timer.start(f'gcn_allreduce_ep{cur_epoch}')
    dist.all_reduce(grad_weight, op=dist.ReduceOp.SUM, group=g_env.world_group)
    torch.cuda.synchronize()
    g_timer.stop(f'gcn_allreduce_ep{cur_epoch}')
    return grad_weight