mnist_hogwild manual breaking of gradient sharing removed (pytorch#138)

ownway22 · May 19, 2017 · 5c41070 · 5c41070
1 parent 89facbe
commit 5c41070
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 5 deletions.
diff --git a/mnist_hogwild/main.py b/mnist_hogwild/main.py
@@ -50,7 +50,7 @@ def forward(self, x):
     torch.manual_seed(args.seed)
 
     model = Net()
-    model.share_memory()
+    model.share_memory() # gradients are allocated lazily, so they are not shared here
 
     processes = []
     for rank in range(args.num_processes):

diff --git a/mnist_hogwild/train.py b/mnist_hogwild/train.py
@@ -7,10 +7,6 @@
 
 def train(rank, args, model):
     torch.manual_seed(args.seed + rank)
-    for param in model.parameters():
-        # Break gradient sharing
-        if param.grad is not None:
-            param.grad.data = param.grad.data.clone()
 
     train_loader = torch.utils.data.DataLoader(
         datasets.MNIST('../data', train=True, download=True,