Updated for class changes