Fix a class I broke