fixed custom class so it's not so ugly :\