Limit training to CPU instead of GPU

[AngshumanChakraborty]

Training on GPU is crashing when loading the CIFAR dataset using
cifar10.prepare(new ProgressBar())
with the following message

ai.djl.engine.EngineException: MXNet engine call failed: MXNetError: cudaMalloc retry failed: out of memory

I have MX130 GPU with 2GB memory on laptop. I have 16GB RAM available for CPU. I want to train on CPU to get around the memory problem. But I cannot select training on CPU even when I have GPU available. How can I train exclusively on CPU even when GPU is available.

[zachgk]

For the Cifar10 dataset, you can change what device it is on by setting a manager for it in the builder:

try(NDManager manager = NDManager.newBaseManager(Device.cpu())) {
    Cifar10 cifar10 = Cifar10.builder()
       .optManager(manager)
       ...
       .build();

    // Train and build within scope of manager
}

The other one to keep in mind is that when you create a new Model, you should create it with the CPU as well like Model.newInstance(modelName, Device.cpu()).

Finally, you will also need to set the device to CPU in your TrainingConfig as well.

Just fix those three things and you should be able to do your training on CPU only.

*Edited to include missed place to update based on comments below

[AngshumanChakraborty]

Thank you, that worked.

Just one thing. I found DefaultTrainingConfig had device as GPU even after setting the model and dataset to CPU. So I changed that too to CPU. But there was no visible change between it having CPU or GPU. So what is the use of optDevices() in DefaultTrainingConfig.

[zachgk]

If we missed setting the device on the DefaultTrainingConfig, it should have actually done the training steps on GPU. It just means you could apparently fit the training batch+backwards on your GPU given the batch size

[AngshumanChakraborty]

Thank you for that.