@jekbradbury I probably figured out what was going wrong. The divisor sample_var.expand_as(input) became zero sometimes, which led to unexpected failure. I forgot to give a lower bound of eps to sample_var. Now the code is running fine. Thanks a lot !! Please have a look at my implementation 159 and tell me how can I improve it.