fixed bug with kl_div