Fixed another little bug in previous commit