search

xinntao / facexlib

FaceXlib aims at providing ready-to-use face-related functions based on current STOA open-source methods.
MIT License
810 stars 143 forks source link

torch.load to correct device. #37

Open jakeh-gc opened 1 year ago

jakeh-gc commented 1 year ago

Use map_location in torch.load to respect user's choice of device.

This fixes a failure when using torch-cpu to load a model with cuda state.

jakeh-gc commented 1 year ago

Without map_location, torch-cpu will throw an exception for model parameters that are expected to be loaded to a cuda device. This means the model.to(device) is never reached.