Hi Frank,
Firstly thank you for your answer. You validate to me, that it was indeed a bug and not something I didn’t understand about autograd!

An easy way to test this is to repeat your computation using double
precision (by using torch.double tensors). If the discrepancy is due
to round-off error it should drop dramatically, say to about 1e-13. If the
discrepancy were to remain around 2e-5, then it would be due to the
two computations not being mathematically equivalent (or to some bug
somewhere).

As I said before the behavior I documented it was because of a nasty little bug. In my tries to make it faster I found an error in my logic (broadcasting is to blame… or me ). After I fixed that, the error drop at 1e-10, even without double tensors!

As an aside, do you ever work with the off-block-diagonal parts of your
tensor X ? If not, you might consider creating, storing, and processing
only the [B, B] -shaped blocks themselves. Given that B << T , doing
so could result in significant time and memory savings.

At least for now, I want all the [T, T] matrix with a bit processing in those blocks for the diagonal. If the logic of the architecture change I would definitely do that to improve my  memory usage. The time savings is already more than enough, after I get rid of that slow for loop!
PS. Thanks again for your help. I’ll mark your answer as the solution, because it was indeed a bug!