Where to set the gpu for loading data

1

when I train the model in a 2 GPUs (e.g. gpu2,gpu3) way, I found undesigned GPU (gpu0) is used to load the data. How can I select the gpu num for loading data?

2

@richard When you create a ndarray you can set the context in which you want the array to be assigned to:
a = mx.nd.ones((1, 10, 10), ctx=mx.gpu(2)) for example.

Please provide a sample code if this isn’t enough to solve your problem.

3

Thanks, in the example below:

for num_batch, data_batch in enumerate(data_train):
module.forward_backward(data_batch)

When checking the context of data_batch before conducting forward_backward processing, I found it is on cpu.
where should I set the context ?