Moving ParamStore Parameters from GPU to CPU

1

Hi,

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.

2

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.
torch.save(param_state, 'param_state.pt')
1 Like