Milestone 2 bug fix