I would recommend to save and load the mode.state_dict(), not the model directly.
That being said, I prefer to push the model to CPU first before saving the state_dict.
This approach makes sure that I’m able to restore the model on all systems, even when no GPU was found.
After loading the model, I use model.to('cuda') to push it to the GPU again.