Fixed broken code after pull request merge