Can you take ‘device‘ as a parameter？

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

Can you take ‘device‘ as a parameter？ #18

Open yizhangliu opened 2 years ago

yizhangliu commented 2 years ago

Thank you for perfect work. But like this: "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')", Sometimes, I will use 'cuda:1'. Please fix it.

Vargol commented 1 year ago

Also there more than just the choice between CUDA and CPU these days, there's MPS too for example

wzxu commented 1 year ago

For myself I manually changed the line to the following to use mps on Apple Silicon Mac.

device = torch.device('mps:0' if torch.backends.mps.is_available() else 'cpu')