Return tuple of loop state and transformed state

[Ian-GL]

If I understood correctly #362 , the new output of the transform should be {%Axon.Loop.State{...}, transformed_state}, so this PR implements said change and updates the tests. This should be the case for other default transforms too? So far I just changed trainer but I guess evaluator should follow suit. Is that right @seanmor5 ?

Closes #362