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).

[kingwingfly]

Hi, bro. How do you find the way to correctly work with multi-gpus on burn.

Right now the batch is moved to the first GPU and users are moving the inputs to the correct device

Could you please tell me how?

I defined Learner with learnerbuilder.devices(vec![Cuda(0), Cuda(1)]), but init model on Cuda(0).

• How to move mini_batch to correct GPU and in which part of the code to do this?
• The model is on GPU(0), but data may be on GPU(1). They are different devices, how can it work?

It's so kind of you if you could provide me any help!

[laggui]

Take a look at the text classification example 🙂 more specifically, how the inputs are moved to the same device as the model.

[kingwingfly]

Thanks a looooot! I managed to. The key point is that the model is forked onto each device and model.devices() will return a Vec<Device> whose first elem is the device the model currenctly on.