Added missing code for one of the comments