Assuming you have precomputed the indices, which would be used for the splitting, you could recreate the order by indexing into the result tensor:
x = torch.tensor([[0.1, 0.3],
                  [0.5, 0.2],
                  [0.7, 0.8],
                  [0.6, 0.6],
                  [0.2, 0.9]])

mlp1_idx = torch.tensor([2, 4])
mlp2_idx = torch.tensor([0, 1, 3,])

x1 = x[mlp1_idx]
x2 = x[mlp2_idx]


# process
...


res = torch.zeros_like(x)
res[torch.cat((mlp1_idx, mlp2_idx), dim=0)] = torch.cat((x1, x2), dim=0)
print(res==x)

Depending on your work flow you might also be able to use e.g. scatter.