Fixed code to work right...