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
require 'cutorch'
require 'cunn'
cutorch.setDevice(1)
cutorch.setDevice(opt.device)
end
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('-batchSize', 50, 'mini-batch size')
cmd:option('-cuda', false, 'Use cuda')
cmd:option('-device', 1, 'Cuda device to use')
cmd:option('-cache', false, 'Cache loaded data.')
cmd:text()