Fix 2D Lorenz with batch_size > 1

I think I fixed what was wrong with the 2D version. I just tried to mimik the behaviour of the code though, so please check very carefully if what I did makes sense!