Hi,
I think your analysis is correct.
The internal implementation is most likely analogous to:
def myloss(A, B, reduction='none'): 
    out = A - B
    return out.pow_(2)

Since we are not allowed to change either A or B inplace. And we want to work in all cases, this leads to the simplest implementation.
Note that the difference with the one that works for your is minimal (only allocating memory slightly differently).
In general, I would advise not to be so close to the max memory. Because, any change to the rest of your code might change the order of the allocation (and the fragementation) and this will start to fail again.