I had similar problem too, that multiplication results had to be in range of 0 < x < 1 because after that, some log operations was resulting NAN, first I was us

I had similar problem too, that multiplication results had to be in range of 0 < x < 1 because after that, some log operations was resulting NAN, first I was using if/else which was slowed down noticeably and after I end up applying sigmoid which speed up things a bit, still I dont know exact cause and I am not brave to dive into that mess again