I already trained without issues by disabling amp for the whole training. But with AMP I got that error. When I’m using AMP all-over, except for that part I mentioned then I’m getting error after 14 iterations. When I’m using AMP all-over, then in the very first iteration. And I’m not using AMP at all, then working fine. So, I was wondering if either normal distribution or kld has any possible connection with AMP behind this problem?