The error you had before the present one, namely

RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling .backward() or autograd.grad() the first time.

is a bit incomplete. Between the two sentences written there, it should read
Think about whether you actually want to take backward of an expression involving gradients and if that is the case specify …
Using anomaly detection mode (with torch.autograd.detect_anomaly()) gives you the line that is causing this: a_tmp[i] = a[i][1].
Indeed this is bad, as the gradient-fn property is only kept per-tensor and not per element, autograd does not know that the backwards of the old iteration isn’t actually needed anymore.
If you change a_tmp to be allocated in each iteration, you won’t have the problem.
So I’m not terribly smart, and as such, I never try this type of clever optimization in the first go of writing a program as you have here with a_tmp and instead optimize for expressivity. This is particularly true here because PyTorch does cache memory allocations, so the overhead of re-allocating is very small (you might use empty to also avoid zeroing if you must). Of course, if you insist being this clever and preallocate, you could also call a_tmp.detach_() in the no-grad block.
Now, smart or not, I have spent quite a while thinking about what it means to write efficient PyTorch code. If you are not contend with the performance of your code, here is one of my rules of thumb that looks applicable: working with lists/loops containing/processing 1-element tensors and very small operands. The things you spell out element by element do look like you would be much happier if you just used a bunch of matmuls or somesuch.
Best regards
Thomas