[FEATURE-REQUEST] Add a --gpus flag

[saattrupdan]

Is your feature request related to a problem? Please describe.
When we have multiple GPUs available, we sometimes want to have control over which ones are used.

Describe the solution you'd like
A flag along the lines of --gpus which could specify a single GPU (--gpus 1 to use the second GPU) or multiple (--gpus 0,1 to use both the first and second GPU). Ideally, some informative logging would be great as well. For instance, "Running on all GPUs", "You're using the first GPU" etc.

Describe alternatives you've considered
An alternative is specifying CUDA_VISIBLE_DEVICES every time the CLI is run, but that's not as intuitive and I often forget it.