Fixing bug. Forgot to push updated class name