pythonpython-3.xmxnet

How to move entire net parameters from gpu to cpu?

My net parameters have been initialized on gpu but I want to move them to cpu.

I've tried copying net as a whole:

net = net.as_in_context(mx.cpu())
# obviously doesn't work

As well as setting individual parameters:

for param in net.collect_params():
    data = net.collect_params()[param].data()
    data = data.as_in_context(mx.cpu())
# does not have any effect either

So how do I move everything to cpu so that subsequent computations use my cpu?

Solution

  • Just realized there is a built-in method for that called reset_ctx:

    for param in net.collect_params().values():
    param.reset_ctx(mx.cpu())