Corrected code based on in class directions