Fixed issues with code

Completed