Improve learner multi-gpu ergonomics

[laggui]

Especially w.r.t data loading, the data should automatically be transferred to the correct device.

Right now the batch is moved to the first GPU and users are moving the inputs to the correct device before inference (which works, but is less nice and might limit multi-gpu usage since the first GPU always gets filled first).