Fix class naming issue