Fix in momentum optimization