Fix issues pointed in comments