#gpu

Two loops run 1000 products of $100 \times 100$ matrices on a GPU and record the Frobenius norm of each result:

# A
log = []
for _ in range(1000):
    Y = X @ W
    log.append(torch.linalg.norm(Y).item())

# B
norms = torch.empty(1000, device="cuda")
for i in range(1000):
    Y = X @ W
    norms[i] = torch.linalg.norm(Y)
log = norms.cpu()

Why is B usually much faster?