Small Changes

Added some code to the end