diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index 8505635..4b1058a 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py
index 98bc2a0..d534ec4 100644
--- a/norch/distributed/distributed.py
+++ b/norch/distributed/distributed.py
@@ -10,10 +10,10 @@ def init_process_group(rank, world_size, backend='nccl'):
     Tensor._C.init_process_group(rank, world_size)
 
 def get_rank():
-    return os.getenv('OMPI_COMM_WORLD_RANK', 0)
+    return int(os.getenv('OMPI_COMM_WORLD_RANK', 0))
 
 def get_world_size():
-    return os.getenv('OMPI_COMM_WORLD_SIZE', 1)
+    return int(os.getenv('OMPI_COMM_WORLD_SIZE', 1))
 
 def broadcast_tensor(tensor, src=0):
 
diff --git a/norch/nn/__pycache__/parameter.cpython-38.pyc b/norch/nn/__pycache__/parameter.cpython-38.pyc
index 14bb069..f85766a 100644
Binary files a/norch/nn/__pycache__/parameter.cpython-38.pyc and b/norch/nn/__pycache__/parameter.cpython-38.pyc differ
diff --git a/norch/nn/parallel.py b/norch/nn/parallel.py
index c021537..338ef18 100644
--- a/norch/nn/parallel.py
+++ b/norch/nn/parallel.py
@@ -23,14 +23,16 @@ class DistributedDataParallel(Module):
         for _, _, parameter in self.parameters():
             dist.broadcast_tensor(parameter)
     
+    @staticmethod
     def allreduce_grads_hook(grad):
         """
         Everytime a gradient is assign to some value, it calculates mean of this gradient among all devices
         """
+        avg_grad = grad
         if isinstance(grad, norch.Tensor):
             dist.allreduce_sum_tensor(grad) 
-            grad /= dist.get_world_size()
-        return grad
+            avg_grad = grad / dist.get_world_size()
+        return avg_grad
     
     def register_grads_hooks(self):
         """
@@ -38,7 +40,6 @@ class DistributedDataParallel(Module):
         """
         for _, _, parameter in self.parameters():
             parameter.register_hook(self.allreduce_grads_hook)
-
     
 
     
diff --git a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc
index 7336b8e..0ee226c 100644
Binary files a/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc and b/norch/optim/optimizers/__pycache__/sgd.cpython-38.pyc differ
diff --git a/norch/optim/optimizers/sgd.py b/norch/optim/optimizers/sgd.py
index 1d785c7..62375f6 100644
--- a/norch/optim/optimizers/sgd.py
+++ b/norch/optim/optimizers/sgd.py
@@ -18,6 +18,7 @@ class SGD(Optimizer):
             velocity = self.momentum * velocity - self.lr * parameter.grad
 
             updated_parameter = parameter + velocity
+            updated_parameter.hooks = parameter.hooks.copy()
 
             setattr(module, name, updated_parameter)
 
diff --git a/train.py b/train.py
index 4eb30c6..666355c 100644
--- a/train.py
+++ b/train.py
@@ -70,16 +70,14 @@ def main():
             loss = criterion(outputs, target)
             
             optimizer.zero_grad()
+            print(f"GRADIENT BEFORE: {model.module.fc2.bias.grad}")
             
-            print("#####################\n\nantes backward")
-            print(model.module.fc1.bias.grad)
             loss.backward()
-            print(model.module.fc1.bias.grad)
-            print("\n\n\n###############\n\npós backward")
+            print(f"GRADIENT AFTER: {model.module.fc2.bias.grad}")
+            print("\n\n")
 
 
             optimizer.step()
-            print("@@")
             break
 
         break