Minor change: updating comments and TODO in model