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?
@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.
Thanks, in the example below:
for num_batch, data_batch in enumerate(data_train):
When checking the context of data_batch before conducting forward_backward processing, I found it is on cpu.
where should I set the context ?