batch-represent: Add device option.

This commit is contained in:
Brandon Amos 2016-06-19 21:43:48 -04:00
parent 85bf6eecc2
commit eaf7a1e19f
2 changed files with 2 additions and 1 deletions

View File

@ -21,7 +21,7 @@ torch.setdefaulttensortype('torch.FloatTensor')
if opt.cuda then if opt.cuda then
require 'cutorch' require 'cutorch'
require 'cunn' require 'cunn'
cutorch.setDevice(1) cutorch.setDevice(opt.device)
end end
opt.manualSeed = 2 opt.manualSeed = 2

View File

@ -24,6 +24,7 @@ function M.parse(arg)
cmd:option('-imgDim', 96, 'Image dimension. nn1=224, nn4=96') cmd:option('-imgDim', 96, 'Image dimension. nn1=224, nn4=96')
cmd:option('-batchSize', 50, 'mini-batch size') cmd:option('-batchSize', 50, 'mini-batch size')
cmd:option('-cuda', false, 'Use cuda') cmd:option('-cuda', false, 'Use cuda')
cmd:option('-device', 1, 'Cuda device to use')
cmd:option('-cache', false, 'Cache loaded data.') cmd:option('-cache', false, 'Cache loaded data.')
cmd:text() cmd:text()