CodeMirror/waifu2x
1
0
Fork 0
mirror of synced 2024-05-16 19:02:21 +12:00
Code Issues Releases Wiki Activity

Fix a bug that batch_size option is not used

Browse source
This commit is contained in:
nagadomi 2018-11-02 18:15:17 +09:00
parent 3654ee41cf
commit 6655d03971
1 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
lib/reconstruct.lua
View file

@ -204,7 +204,7 @@ function reconstruct.scale_rgb(model, scale, x, offset, block_size, batch_size)
collectgarbage()
return output
end
function reconstruct.image(model, x, block_size)
function reconstruct.image(model, x, block_size, batch_size)
local i2rgb = false
if x:size(1) == 1 then
local new_x = torch.Tensor(3, x:size(2), x:size(3))
@ -216,17 +216,17 @@ function reconstruct.image(model, x, block_size)
end
if reconstruct.is_rgb(model) then
x = reconstruct.image_rgb(model, x,
reconstruct.offset_size(model), block_size)
reconstruct.offset_size(model), block_size, batch_size)
else
x = reconstruct.image_y(model, x,
reconstruct.offset_size(model), block_size)
reconstruct.offset_size(model), block_size, batch_size)
end
if i2rgb then
x = image.rgb2y(x)
end
return x
end
function reconstruct.scale(model, scale, x, block_size)
function reconstruct.scale(model, scale, x, block_size, batch_size)
local i2rgb = false
if x:size(1) == 1 then
local new_x = torch.Tensor(3, x:size(2), x:size(3))
@ -239,11 +239,11 @@ function reconstruct.scale(model, scale, x, block_size)
if reconstruct.is_rgb(model) then
x = reconstruct.scale_rgb(model, scale, x,
reconstruct.offset_size(model),
block_size)
block_size, batch_size)
else
x = reconstruct.scale_y(model, scale, x,
reconstruct.offset_size(model),
block_size)
block_size, batch_size)
end
if i2rgb then
x = image.rgb2y(x)