(Minor) Fix in logic