You sir are fantastic! This worked and makes perfect sense.
Not surprisingly I immediately ran into another question (which there is probably good documentation for somewhere).
Initially I got
RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.

Which I then followed and restarted training. For each LBFGS internal iteration it takes a longer and longer time.
Is the graph growing and growing with each iteration? More generally what would you suggest I read / search for to better understand what is happening here.
Thanks again for your solution to my problem and the great library!