I see your loop as
for i in range(n):
  ba_k[i] = f1(ab,ba,i)
  ab_k[i] = f2(ab,ba,i)

so it looks parallelizable (non-recurrent) to me… may be too complicated though…



 legoh:

I probably could just create a tensor and keep stacking them?


You don’t have to create tensors, as (m_ab - torch.max(…)) expressions already allocate memory for partial results. You just keep them around in a list, and stack once after a loop.