浏览代码

Fix NaN bug in iproc.white_noise

nagadomi 9 年之前
父节点
当前提交
28cc864c0c
共有 1 个文件被更改,包括 4 次插入1 次删除
  1. 4 1
      lib/iproc.lua

+ 4 - 1
lib/iproc.lua

@@ -99,7 +99,10 @@ function iproc.white_noise(src, std, rgb_weights, gamma)
 
    local dest
    if gamma ~= 0 then
-      dest = src:clone():pow(gamma):add(noise):pow(1.0 / gamma)
+      dest = src:clone():pow(gamma):add(noise)
+      dest[torch.lt(dest, 0.0)] = 0.0
+      dest[torch.gt(dest, 1.0)] = 1.0
+      dest:pow(1.0 / gamma)
    else
       dest = src + noise
    end