Fix bug for nan checks