Updated comments to use the correct imports