DPT Synchronize() performance fix. GPU flag fix

soumith · Feb 10, 2016 · edcbd57 · edcbd57
1 parent c73fb73
commit edcbd57
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/train.lua b/train.lua
@@ -163,7 +163,10 @@ function trainBatch(inputsCPU, labelsCPU)
    optim.sgd(feval, parameters, optimState)
 
    -- DataParallelTable's syncParameters
-   model:apply(function(m) if m.syncParameters then m:syncParameters() end end)
+   if model.needsSync then
+      model:syncParameters()
+   end
+
 
    cutorch.synchronize()
    batchNumber = batchNumber + 1

diff --git a/util.lua b/util.lua
@@ -10,8 +10,9 @@ function makeDataParallel(model, nGPU)
          cutorch.setDevice(i)
          model:add(model_single:clone():cuda(), i)
       end
-      cutorch.setDevice(opt.GPU)
    end
+   cutorch.setDevice(opt.GPU)
+
    return model
 end