When I call model.to(‘cpu’) on a model that’s currently in gpu, does it do anything aside from move the individual tensors in the state_dict to cpu? I’m trying to move the tensors individually because I need a portion of the parameters to remain in gpu because they’re shared by another model that’s still running. I’ve tried the following to edit the state_dict and then load it back:
updated = model2.state_dict()
for i in updated.keys():
print(f’before: {updated[i]}’)
updated[i] = updated[i].to(‘cpu’)
print(f’after: {updated[i]}’)
model2.load_state_dict(updated)
print(model2.state_dict())
I can see that the tensors in “updated” are moving from cuda to cpu in the before and after calls, but when I load it back into model2, all elements of the state_dict still appear to be in cuda. How do I move them individually so that I can then filter some keys out?
Thanks!