When you torch.load, it tries to put the tensors back on the device they were saved from
[torch.load()](https://pytorch.org/docs/stable/generated/torch.load.html#torch.load) uses Python’s unpickling facilities but treats storages, which underlie tensors, specially. They are first deserialized on the CPU and are then moved to the device they were saved from. If this fails (e.g. because the run time system doesn’t have certain devices), an exception is raised. However, storages can be dynamically remapped to an alternative set of devices using the map_location argument.

This ends up with an error if you (e.g.) try to load something on cpu after it was saved on gpu, so everywhere we call torch.load we map to cpu. I just missed that piece in my original PR with these changes.