made some changes based on github comments