diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 01ab4b3..00a1896 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/tensor.o b/build/tensor.o
index 4246623..17fca73 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index 362e0a2..20e600b 100644
--- a/norch/csrc/cuda.cu
+++ b/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); 
diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp
index bb753b2..0a1cc69 100644
--- a/norch/csrc/distributed.cpp
+++ b/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));
 
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index b1c9408..cc6b922 100644
--- a/norch/csrc/tensor.cpp
+++ b/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);
         }
     }
 
diff --git a/norch/libtensor.so b/norch/libtensor.so
index dc9871e..f50c17c 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ