Can you take ‘device‘ as a parameter？

[yizhangliu]

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]

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

[wzxu]

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')