Training: Add cuda device option.

cmusatyalab · Jun 14, 2016 · 92684ab · 92684ab
1 parent a1b3251
commit 92684ab
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/training/main.lua b/training/main.lua
@@ -14,7 +14,7 @@ print(opt)
 
 if opt.cuda then
    require 'cutorch'
-   cutorch.setDevice(1)
+   cutorch.setDevice(opt.device)
 end
 
 torch.save(paths.concat(opt.save, 'opts.t7'), opt, 'ascii')

diff --git a/training/opts.lua b/training/opts.lua
@@ -27,6 +27,7 @@ function M.parse(arg)
               'Home of dataset. Split into "train" and "val" directories that separate images by class.')
    cmd:option('-manualSeed', 2, 'Manually set RNG seed')
    cmd:option('-cuda', true, 'Use cuda.')
+   cmd:option('-device', 1, 'Cuda device to use.')
    cmd:option('-cudnn', true, 'Convert the model to cudnn.')
 
    ------------- Data options ------------------------