Training: Add cuda device option.
This commit is contained in:
parent
a1b3251f96
commit
92684ab771
|
@ -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')
|
||||
|
|
|
@ -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 ------------------------
|
||||
|
|
Loading…
Reference in New Issue