diff --git a/norch/csrc/distributed.cpp b/norch/csrc/distributed.cpp
index cf84b85..5239eaa 100644
--- a/norch/csrc/distributed.cpp
+++ b/norch/csrc/distributed.cpp
@@ -63,8 +63,8 @@ void allreduce_mean_tensor(Tensor* tensor) {
     // Perform NCCL AllReduce operation to calculate the mean of all tensors across all processes
     NCCL_CHECK(ncclAllReduce(tensor->data, tensor->data, tensor->size, ncclFloat, ncclSum, nccl_comm, stream));
 
-    tensor_div_scalar(tensor, tensor->data);
-    
+    tensor_div_scalar_cuda(tensor, world_size, tensor->data);
+
     cudaStreamSynchronize(stream);
     cudaStreamDestroy(stream);
 }
diff --git a/norch/csrc/distributed.h b/norch/csrc/distributed.h
index 1992b8f..3a94983 100644
--- a/norch/csrc/distributed.h
+++ b/norch/csrc/distributed.h
@@ -24,6 +24,7 @@ extern "C" {
   void init_process_group(int rank, int world_size);
   void broadcast_tensor(Tensor* tensor);
   void allreduce_sum_tensor(Tensor* tensor);
+  void allreduce_mean_tensor(Tensor* tensor);
 }
 
 #endif /* DISTRIBUTED_H */
diff --git a/norch/nn/__init__.py b/norch/nn/__init__.py
index 8f7997b..560eef8 100644
--- a/norch/nn/__init__.py
+++ b/norch/nn/__init__.py
@@ -1,4 +1,5 @@
 from .modules import *
 from .activation import *
 from .loss import *
-from .functional import *
\ No newline at end of file
+from .functional import *
+from .distributed import *
\ No newline at end of file
diff --git a/norch/nn/distributed.py b/norch/nn/distributed.py
index a9a5de5..7d5060b 100644
--- a/norch/nn/distributed.py
+++ b/norch/nn/distributed.py
@@ -14,7 +14,7 @@ class DistributedDataParallel(Module):
     def backward(self):
         self.module.backward()
         for module, name, parameter in self.parameters():
-            dist.allreduce_sum_tensor(parameter)
+            dist.allreduce_mean_tensor(parameter)