From eaf7a1e19f7b919c43fc19d7d5227e3827f4f0f0 Mon Sep 17 00:00:00 2001 From: Brandon Amos Date: Sun, 19 Jun 2016 21:43:48 -0400 Subject: [PATCH] batch-represent: Add device option. --- batch-represent/main.lua | 2 +- batch-represent/opts.lua | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/batch-represent/main.lua b/batch-represent/main.lua index 28c8faf..a2fdb13 100755 --- a/batch-represent/main.lua +++ b/batch-represent/main.lua @@ -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 diff --git a/batch-represent/opts.lua b/batch-represent/opts.lua index 2813d70..7841c28 100644 --- a/batch-represent/opts.lua +++ b/batch-represent/opts.lua @@ -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()