for some reason, this fixes the nans and error in the test