Allow specifying device to use when training

We should expose a way to mutate this variable to select a specific GPU to use.