From 92684ab771f028da1036545b685fe7d02cf7745f Mon Sep 17 00:00:00 2001 From: Brandon Amos Date: Tue, 14 Jun 2016 15:55:46 -0400 Subject: [PATCH] Training: Add cuda device option. --- training/main.lua | 2 +- training/opts.lua | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/training/main.lua b/training/main.lua index 5df5c9d..104bd5a 100755 --- 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 index 9089f85..20f67a0 100644 --- 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 ------------------------