Bug fix (better if statement) (#306)