CodeMirror/waifu2x
1
0
Fork 0
mirror of synced 2024-05-19 20:32:22 +12:00
Code Issues Releases Wiki Activity

Add -gpu option in train.lua

Browse source
This commit is contained in:
nagadomi 2015-11-13 19:26:58 +09:00
parent c345f35ed0
commit 42bd89151e
2 changed files with 4 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
lib/settings.lua
View file

@ -17,6 +17,7 @@ local cmd = torch.CmdLine()
cmd:text()
cmd:text("waifu2x-training")
cmd:text("Options:")
cmd:option("-gpu", -1, 'GPU Device ID')
cmd:option("-seed", 11, 'RNG seed')
cmd:option("-data_dir", "./data", 'path to data directory')
cmd:option("-backend", "cunn", '(cunn|cudnn)')

3
train.lua
View file

@ -203,6 +203,9 @@ local function train()
collectgarbage()
end
end
if settings.gpu > 0 then
cutorch.setDevice(settings.gpu)
end
torch.manualSeed(settings.seed)
cutorch.manualSeed(settings.seed)
print(settings)