Address github comments