Task2 bug fix