fix distributed ncclCommInitRank

norch/csrc/cuda.cu

@@ -13,7 +13,7 @@ __host__ void cpu_to_cuda(Tensor* tensor, int device_id) {
     cudaGetDeviceCount(&deviceCount);
     if (device_id >= deviceCount) {
         fprintf(stderr, "Could not send tensor to device %d, only %d devices available\n", device_id, deviceCount);
-            exit(1);
+        exit(1);
     }
     
     cudaSetDevice(device_id); 

norch/csrc/distributed.cpp

@@ -30,6 +30,8 @@ void init_process_group(int env_rank, int env_world_size) {
     // send nccl unique id to all processes
     MPI_CHECK(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD));
 
+    cudaSetDevice(rank); 
+
     // init NCCL communication group
     NCCL_CHECK(ncclCommInitRank(&nccl_comm, world_size, nccl_id, rank));
 

norch/csrc/tensor.cpp

@@ -83,7 +83,8 @@ extern "C" {
         }
 
         else {
-            printf("Could not send tensor to device %d", device_id);
+            fprintf(stderr, "Could not send tensor to device %d", device_id);
+            exit(1);
         }
     }