Added small changes to assignment