pytorch
diff --git a/‎test/cpp/c10d/ProcessGroupNCCLErrorsTest.cpp
Copy file name to clipboardExpand all lines: test/cpp/c10d/ProcessGroupNCCLErrorsTest.cpp
+15-2Lines changed: 15 additions & 2 deletions b/‎test/cpp/c10d/ProcessGroupNCCLErrorsTest.cpp
Copy file name to clipboardExpand all lines: test/cpp/c10d/ProcessGroupNCCLErrorsTest.cpp
+15-2Lines changed: 15 additions & 2 deletions
@@ -254,12 +254,16 @@ class ProcessGroupNCCLErrorsTest : public ::testing::Test {
   void SetUp() override {
     // Enable LOG(INFO) messages.
     c10::initLogging();
-    size_t numDevices = 1; // One device per rank (thread)
+    size_t numDevices = cudaNumDevices();
     TemporaryFile file;
     store_ = c10::make_intrusive<::c10d::FileStore>(file.path, 1);
 
+    at::cuda::OptionalCUDAGuard deviceGuard;
     tensors_.resize(numDevices);
-    tensors_[0] = at::empty({3, 3}, at::kCUDA);
+    for (const auto i : c10::irange(numDevices)) {
+      deviceGuard.set_index(i);
+      tensors_[i] = at::ones({3, 3}, at::kCUDA);
+    }
   }
 
   void TearDown() override {
@@ -282,6 +286,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLErrorsBlocking) {
 
   auto work = pg.allreduce(tensors_);
   work->wait();
+  EXPECT_TRUE(work->isSuccess());
   EXPECT_EQ(1, pg.getNCCLCommCacheSize());
 
   // Now run all reduce with errors.
@@ -291,6 +296,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLErrorsBlocking) {
 
   // Verify the work item failed.
   EXPECT_TRUE(work->isCompleted());
+  EXPECT_FALSE(work->isSuccess());
   EXPECT_THROW(work->wait(), std::runtime_error);
 
   // Communicators might be aborted here, further operations would fail.
@@ -308,6 +314,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLTimedoutErrorsBlocking) {
 
   auto work = pg.allreduce(tensors_);
   work->wait();
+  EXPECT_TRUE(work->isSuccess());
   EXPECT_EQ(1, pg.getNCCLCommCacheSize());
 
   // Now run all reduce with errors.
@@ -329,6 +336,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLErrorsNonBlocking) {
 
   auto work = pg.allreduce(tensors_);
   pg.barrier()->wait();
+  EXPECT_TRUE(work->isSuccess());
   EXPECT_EQ(1, pg.getNCCLCommCacheSize());
 
   // Now run all reduce with errors.
@@ -339,7 +347,10 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLErrorsNonBlocking) {
   work->wait();
   pg.barrier()->wait();
 
+  // Verify the work item failed.
   EXPECT_TRUE(work->isCompleted());
+  EXPECT_FALSE(work->isSuccess());
+
   // Communicators might be aborted here, further operations would fail.
 }
 
@@ -415,6 +426,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLErrorsNoHeartbeat) {
   // Normal collective case.
   auto work = pg.allreduce(tensors_);
   work->wait();
+  EXPECT_TRUE(work->isSuccess());
 
   work = pg.allreduce(tensors_);
   {
@@ -428,6 +440,7 @@ TEST_F(ProcessGroupNCCLErrorsTest, testNCCLErrorsNoHeartbeat) {
     EXPECT_TRUE(pg.getErrorCaughtFlag());
   }
   work->wait();
+  EXPECT_TRUE(work->isSuccess());
   EXPECT_TRUE(traces.size() > 0);
   auto filename = c10::str(tempFilename, 0);
   auto traceFromStorage = readTraceFromFile(filename, traces.size());