Bugfix: use correct class