Moving ParamStore Parameters from GPU to CPU


After training a model on the GPU, I’m having trouble moving the parameters in the Pyro ParamStore from the GPU to the CPU (since storage/deploy/predict is going to be done on the CPU). Is there an easy way to move ParamStore parameters from the GPU to CPU (or vice versa)? I couldn’t find any method to easily do this.

Ok, after tinkering around a bit, the method I came up with is below. Not sure if this is the best way to do it, but just thought I’d share. If anyone has a better way, feel free to let me know.

# Store relevant param store attributes in dict.
param_state = pyro.get_param_store().get_state()

# Convert tensors from cuda to cpu.
with torch.no_grad():
    for name, value in param_state['params'].items():
        param_state['params'][name] = value.cpu()

# Save the param store dict., '')
1 Like