set the right cuda device in DataLoadingThread (pytorch#1742)

Summary:
Pull Request resolved: pytorch#1742

We need to set the cuda device in DataLoadingThread.

Without it, GPU 0 will be used by the data loading threads of other GPUs.

Reviewed By: joshuadeng

Differential Revision: D54402994

fbshipit-source-id: 47a071ac40eebf3703e3a63fffdd256933b5b12c

torchrec/distributed/train_pipeline.py

@@ -78,6 +78,10 @@ def __init__(
         self._buffer_empty_event.set()
 
     def run(self) -> None:
+        if self._device.type == "cuda" and torch.cuda.is_available():
+            # set the current device the same as the one used in the main thread
+            torch.cuda.set_device(self._device)
+
         while not self._stop:
             self._buffer_empty_event.wait()
             # Set the filled event to unblock progress() and return.