From 0359ec765e71b7ff1d6567d0d469f5638744c689 Mon Sep 17 00:00:00 2001
From: "Gao, Xiang" <qasdfgtyuiop@gmail.com>
Date: Thu, 10 Jun 2021 15:52:46 -0700
Subject: [PATCH] Set to the correct device in ProcessGroupUCC::initComm

---
 src/torch_ucc.cpp | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/torch_ucc.cpp b/src/torch_ucc.cpp
index 941fb06..55ae70d 100644
--- a/src/torch_ucc.cpp
+++ b/src/torch_ucc.cpp
@@ -802,6 +802,11 @@ c10::intrusive_ptr<ProcessGroup> ProcessGroupUCC::createProcessGroupUCC(
 
 void ProcessGroupUCC::initComm(c10::Device dev) {
   if (!comm) {
+#ifdef USE_CUDA
+    if (dev.is_cuda()) {
+      c10::cuda::set_device(dev.index());
+    }
+#endif
     comm = CommPG::get_comm(comm_id, dev, &oob);
     comm->ucx_connect_eps(eps, &oob);
     comm->ucc_create_team(team, &oob);