Is lu_unpack (src 3) as efficient as possible? For instance let’s have a look at this part:
P = torch.eye(sz, device=LU_data.device, dtype=LU_data.dtype)
final_order = list(range(sz))
for k, j, in enumerate(LU_pivots_zero_idx):
    final_order[k], final_order[j] = final_order[j], final_order[k]
P = P.index_select(1, torch.as_tensor(final_order, device=LU_pivots.device))

Why isn’t final_order allocated on LU_pivots.device from the start? Doesn’t final_order[j] forces a sync if LU_pivots is on a GPU?