Modified code for assignment