Optimise device allocation for all attributes used in the training loop

This MR optimises device allocation for all model parameters participating in the training loop.

It closes #78 (closed).

No significant effect was observed in terms of processing speed, however we still explicitly move transforms and losses as a cautionary measure.