Bug fix so code runs