Updated code from lab