[CI] linux://doc:datasets_train is failing/flaky on master.

[rickyyx]

• 6bac4ab FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 34cd1cb FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 5229df1 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 44fc6a3 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• c9875ed FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 60840b5 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 297341e FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 217dcd7 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 9a7aa24 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 47aa6c4 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• da6adbb FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 081ce2f FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• d7b2b49 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 3ede830 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 65a8b0f FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 91040d9 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 8486ddd FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 7e939d7 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 4732753 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 68a35fc FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)
• 23b3a59 FAILED Buildkite 📺 :database: 🚂 Datasets Train Integration GPU Tests and Examples (Python 3.7)

....
Generated from flaky test tracker. Please do not edit the signature in this section.
DataCaseName-linux://doc:datasets_train-END
....

[rickyyx]

Failure on release branch: https://buildkite.com/ray-project/oss-ci-build-branch/builds/488#0183cd54-7535-4e9e-8420-3076ccebc871

[matthewdeng]

Hmmm @amogkam any ideas about this one?

opt/bin/doc/datasets_train.runfiles/com_github_ray_project_ray/doc/source/ray-core/_examples/datasets_train/datasets_train.py", line 350, in train_epoch
--
  | loss = criterion(outputs, labels.float())
  | File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
  | return forward_call(*input, **kwargs)
  | File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/modules/loss.py", line 716, in forward
  | reduction=self.reduction)
  | File "/opt/miniconda/lib/python3.7/site-packages/torch/nn/functional.py", line 2960, in binary_cross_entropy_with_logits
  | return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
  | RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!

I originally thought it was due to #29104 but see that it's already checked in...

[amogkam]

The example should use train.torch.get_device() instead of manually doing it with local rank

# Setup device.
    device = torch.device(
        f"cuda:{session.get_local_rank()}"
        if use_gpu and torch.cuda.is_available()
        else "cpu"
    )

[amogkam]

#29305