DeepSpeed

135ffe30 - fix some issue with saving the variance for normalize forward kernel