Fix nan propagation in tf.relu (tensorflow#1001)

BUG Default behavior changed with newest NVidia Drivers in Mac OS >=10.13.4 so we have to check for NaN explicitly.
indutny · Apr 26, 2018 · d3bf768 · d3bf768
1 parent 24c78cc
commit d3bf768
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/kernels/webgl/unaryop_gpu.ts b/src/kernels/webgl/unaryop_gpu.ts
@@ -46,7 +46,9 @@ const CHECK_NAN_SNIPPET = `if (isNaN(x)) return x;`;
 
 export const ABS = `return abs(x);`;
 
-export const RELU = `return (x < 0.0) ? 0.0 : x;`;
+export const RELU = CHECK_NAN_SNIPPET + `
+  return (x < 0.0) ? 0.0 : x;
+`;
 
 export const ELU = `return (x >= 0.0) ? x : (exp(x) - 1.0);`;