Just want to add something here:
I ran into the same problem when I resumed a training.

“When you call torch.load() on a file which contains GPU tensors, those tensors will be loaded to GPU by default.”

In my case, I run trainings on different GPUs on a multi GPU-system. So from time to time my GPU changes.
So, to avoid this loading issues:

torch.load(checkpoint_file, map_location=‘cpu’)

Later you can push everything on the right device.