diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp
index 666c435..8124ec0 100644
--- a/norch/csrc/distributed.cpp
+++ b/norch/csrc/distributed.cpp
@@ -46,17 +46,13 @@ void broadcast_tensor(Tensor* tensor) {
     cudaStreamDestroy(stream);
 }
 
-void allreduce_mean_tensor(Tensor* tensor) {
+void allreduce_sum_tensor(Tensor* tensor) {
     cudaStream_t stream;
     cudaStreamCreate(&stream);
 
     // Perform NCCL AllReduce operation to calculate the sum of all tensors across all processes
     NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
 
-    // Divide the sum by the world size to get the mean
-    float scale = 1.0f / world_size;
-    scalar_mul_tensor_cuda(tensor, scale, tensor->data);
-
     cudaStreamSynchronize(stream);
     cudaStreamDestroy(stream);
 }
diff --git a/norch/csrc/distributed.h b/norch/csrc/distributed.h
index 40da266..1040fc0 100644
--- a/norch/csrc/distributed.h
+++ b/norch/csrc/distributed.h
@@ -21,6 +21,8 @@
 
 extern "C" {
   void init_process_group(int rank, int world_size);
+  void broadcast_tensor(Tensor* tensor);
+  void allreduce_sum_tensor(Tensor* tensor);
 }
 
 #endif /* DISTRIBUTED_H */
diff --git a/norch/csrc/distributed.o b/norch/csrc/distributed.o
deleted file mode 100755
index 144f330..0000000
Binary files a/norch/csrc/distributed.o and /dev/null differ
diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py
index 3d5574e..d033ede 100644
--- a/norch/distributed/distributed.py
+++ b/norch/distributed/distributed.py
@@ -16,7 +16,7 @@ def broadcast_tensor(tensor):
     
     Tensor._C.broadcast_tensor(tensor)
 
-def allreduce_mean_tensor(tensor):
+def allreduce_sum_tensor(tensor):
 
     Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
     Tensor._C.allreduce_mean_tensor.restype = None
diff --git a/train.py b/train.py
index 6cea436..8f07a41 100644
--- a/train.py
+++ b/train.py
@@ -10,17 +10,22 @@ def main():
 
     dist.init_process_group(rank, world_size)
 
-    tensor = norch.Tensor([1,2,3])
-    tensor = rank * tensor
-    print(f"BEFORE on rank {rank}: ", tensor, '\n\n')
+    tensor = norch.Tensor([1,1,1])
+    tensor = (rank + 1) * tensor
+    print(f"BEFORE on rank {rank}: {tensor} \n\n")
+
+    dist.allreduce_sum_tensor(tensor)
+
+    print(f"AFTER ALLREDUCE on rank {rank}: {tensor} \n\n")
+
+    print("###############\n\n\n")
+
+    tensor = tensor * 10
+    print(f"BEFORE BROADCAST on rank {rank}: {tensor} \n\n")
 
     dist.broadcast_tensor(tensor)
 
-    print(f"AFTER BROADCAST on rank {rank}: ", tensor, '\n\n')
-
-    dist.allreduce_mean_tensor(tensor)
-
-    print(f"AFTER ALLREDUCE MEAN on rank {rank}: ", tensor, '\n\n')
+    print(f"AFTER BROADCAST on rank {rank}: {tensor} \n\n")
 
 if __name__ == "__main__":
     main()