Different devices for different models

[Huan80805]

Hi, all.
I'm using three different models A,B,C in my project.
In short, I'm trying to fine-tune A to produce prompts, the prompts is then used to prompt B for response, then B's response is scored by C, and finally the score is used to fine-tune A with RL.
I would like to put A in device_1, and put B,C on device_2.
I am able to put A on device_1 and C on device_2. But as long as I try to load B on device_2, the error pops. (putting both A and B on device_1, C on device_2 works well)
return F.linear(input, self.weight, self.bias). RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
By checking, the input (A's input) is on device_2 and A's weight is on device_1.
Please enlighten me on how to distribute different models to different devices in accelerate framework.