Thanks for sharing. So if I understand correctly, the torch.cat is a bandwidth-bound problem rather than a compute-bound problem.
I also tried a larger model. You’re right! GPU is faster than CPU.
import time
import torch
import torch.nn as nn
import torch.autograd as autograd


class netG(nn.Module):
    def __init__(self):
        super(netG, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(32, 1024, bias=False),
            nn.Linear(1024, 512, bias=True)
        )

    def forward(self, x):
        return self.net(x)



class netD(nn.Module):
    def __init__(self):
        super(netD, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(512, 512),
            nn.Linear(512, 1)
        )

    def forward(self, x):
        return self.net(x)


if __name__ == '__main__':
    device = torch.device('cuda:0') if torch.cuda.is_available() else 'cpu'
    # device = 'cpu'
    print(device)
    torch.cuda.init()
    z = torch.randn(32, device=device)
    G = netG().to(device)
    D = netD().to(device)
    loss = D(G(z))
    if device == torch.device('cuda:0'):
        torch.cuda.synchronize(device=device)
    start = time.time()
    grad_g = autograd.grad(loss, list(G.parameters()), create_graph=True, retain_graph=True)
    grad_d = autograd.grad(loss, list(D.parameters()), create_graph=True, retain_graph=True)
    if device == torch.device('cuda:0'):
        torch.cuda.synchronize(device=device)
    end = time.time()
    grad_time = start - end
    print('time cost for autograd: {:.5f}'.format(grad_time))

    if device == torch.device('cuda:0'):
        torch.cuda.synchronize(device=device)
    start = time.time()
    grad_g_vec = torch.cat([g.contiguous().view(-1) for g in grad_g])
    grad_d_vec = torch.cat([g.contiguous().view(-1) for g in grad_d])
    if device == torch.device('cuda:0'):
        torch.cuda.synchronize(device=device)
    end = time.time()
    grad_time = start - end
    print('time cost for cat: {:.5f}'.format(grad_time))

Result on big model
cpu
time cost for autograd: -0.03366
time cost for cat: -0.00306
cuda:0
time cost for autograd: -0.00318
time cost for cat: -0.00147