Here is my new test_lr_batch function
def test_lr_batch():
    def check_equal(x, y, eps=1e-8):
        assert torch.sum(torch.abs(x - y)) <= eps

    tt1, tt2 = 0, 0
    x = torch.randn(4, 4800, 160, dtype=torch.float64).cuda()
    beta = torch.randn(4, 160, 1, dtype=torch.float64).cuda()
    y = torch.bmm(x, beta).squeeze(2)
    beta = beta.squeeze(2)
    torch.cuda.synchronize()
    for i in range(200):
        torch.cuda.synchronize()
        t1 = time.time()

        beta2 = lr_batch_fit_c(x, y)

        torch.cuda.synchronize()
        t2 = time.time()

        beta3 = lr_batch_fit_py(x, y)

        torch.cuda.synchronize()
        t3 = time.time()
        tt1 += (t2 - t1)
        tt2 += (t3 - t2)
        check_equal(beta, beta2)
        check_equal(beta, beta3)
    print("test_lr_batch  c_time %.3fs py_time %.3fs" % (tt1, tt2))

Based on the new function, the cpp version is still slower than the py version.
Is the new function fine for testing, or should I also put similar stuff ( torch.cuda.synchronize() ) in the cpp code?