[train][docs] update docstrings/quickstarts to work when use_gpu=True

[matthewdeng]

Signed-off-by: Matthew Deng [email protected]

Fixes Trainer docstrings and quickstarts to work when use_gpu=True.

Why are these changes needed?

1. Updated iter_torch_batches to move to the proper device.
   1. This is needed for both torch and horovod which use iter_torch_batches.
   2. This is not needed for tensorlfow and huggingface which handle device transfer natively.
2. Extracted use_gpu to the top of each code snippet for easier control handling.

Related issue number

Closes #31684

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[Yard1]

Orthogonal to this, but I think that if we detect we are in a train session, iter_torch_batches should automatically use the default device unless specified otherwise.

[maxpumperla]

dumb, yet effective trick for cases you want to show users you can use gpus, but want cpus on CI:

use_gpu = True  # include in docs

use_gpu = False # exclude

<code using use_gpu>

[matthewdeng]

Orthogonal to this, but I think that if we detect we are in a train session, iter_torch_batches should automatically use the default device unless specified otherwise.

@Yard1 yeah I was thinking the same thing, cc @stephanie-wang this could be a nice extension to the DatasetIterator!

[Yard1]

I've got a PR here, comments appreciated - #31745

[matthewdeng]

@amogkam let me know how you want to coordinate this with #31753. If you're able to get your changes in by 2.3 I can update this PR to not set device.

[amogkam]

Ah sorry- let's get this in ASAP...let me know when it's ready to merge.

I can update this documentation in that PR.

[matthewdeng]

@amogkam I think it's good to merge as is!