Fix a bug brought on by refactoring