albanD:

We need to go in cpp, discover the work that needs to be done, start GPU kernels for everything.
That takes some time.


I modified the code like so to investigate your explanantion:
for i, l in enumerate(losses):
    t = time()
    # just to do something with the gradients
    model.zero_grad()
    l.backward(retain_graph=True)
    print(i, time() - t)

What I get from this is
...
24 0.0007534027099609375
25 0.0007722377777099609
26 0.0007855892181396484
27 0.0024886131286621094
28 0.0009698867797851562
29 0.0012204647064208984
30 0.0010230541229248047
31 0.0008389949798583984
32 0.0024178028106689453
33 0.0007109642028808594
34 0.0007128715515136719
35 0.0006940364837646484
36 0.0009226799011230469
37 0.0046350955963134766
38 0.004754304885864258
39 0.004813194274902344
40 0.004805564880371094
41 0.004813432693481445
...

Why do you think does it take four times as long to go to cpp in iteration 37 compared with iteration 36?
Here’s an infographics over a bunch of runs on my dedicated nv 2080 (sampling rate ~ 1 / (time() - t), iteration: that’s the index into the batch in the loop across losses):

However, thanks for the pointer to autograd hacks; I’ll try that out.