The map_location argument remaps the tensors to the specified device. This is useful, if your current machine does not have a GPU and you are trying to load a tensor, which was saved as a CUDATensor.

From the docs:

 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.


You could push the model to the GPU before loading the state_dict. However, the tensors have to be pushed somewhere to the device, so I guess you won’t save any time.