I actually figured out what the issue was. I am appending the RNG state to the models state dict when saving a checkpoint. So when loading from a checkpoint, the map_location was sending it to the GPU. Silly oversight on my part.
Thanks for responding anyway though! the state tensor being a parameter would have been a good explanation